Upload Caduceus

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5cf9b060f9fe236a72eeb4df44d905816bbf3ce013462524edcfbb7272b8c947
-size 2174536
+oid sha256:16d3acb10a57ce482dd0799e59fd8616b83ce414143b04d68d95a9ab8cd8180e
+size 2173880

modeling_caduceus.py

@@ -158,7 +158,7 @@ class CaduceusMixerModel(nn.Module):
         self.rcps = config.rcps
         self.residual_in_fp32 = config.residual_in_fp32
 
-        self.embeddings = torch.compile(CaduceusEmbeddings(config, **factory_kwargs))
+        self.embeddings = CaduceusEmbeddings(config, **factory_kwargs)
 
         # Mamba changes the order of residual and layer norm:
         # Instead of LN -> Attn / MLP -> Add, we do:
@@ -377,12 +377,12 @@ class CaduceusForMaskedLM(CaduceusPreTrainedModel):
         factory_kwargs = {"device": device, "dtype": dtype}
         self.caduceus = Caduceus(config, **factory_kwargs, **kwargs)
         if config.rcps:
-            self.lm_head = torch.compile(RCPSLMHead(
+            self.lm_head = RCPSLMHead(
                 complement_map=self.config.complement_map,  # Use caduceus config as it might have been updated
                 vocab_size=self.config.vocab_size,  # Use caduceus config as it might have been updated
                 true_dim=config.d_model,
                 dtype=dtype
-            ))
+            )
         else:
             self.lm_head = nn.Linear(
                 config.d_model,

modeling_rcps.py

@@ -144,7 +144,7 @@ class RCPSMambaBlock(nn.Module):
         super().__init__()
         self.residual_in_fp32 = residual_in_fp32
         self.fused_add_norm = fused_add_norm
-        self.mixer = torch.compile(RCPSWrapper(mixer_cls(dim)))
+        self.mixer = RCPSWrapper(mixer_cls(dim))
         norm_f = norm_cls(dim)
         self.norm = norm_f if fused_add_norm else RCPSAddNormWrapper(norm_f)