Update modeling_neuroclr.py

modeling_neuroclr.py

@@ -56,6 +56,22 @@ class NeuroCLR(nn.Module):
 
         return h, z
 
+class NeuroCLRModel(PreTrainedModel):
+    """
+    Loads with:
+      AutoModel.from_pretrained(..., trust_remote_code=True)
+    """
+    config_class = NeuroCLRConfig
+    base_model_prefix = "neuroclr"
+
+    def __init__(self, config: NeuroCLRConfig):
+        super().__init__(config)
+        self.neuroclr = NeuroCLR(config)
+        self.post_init()
+
+    def forward(self, x: torch.Tensor, **kwargs):
+        h, z = self.neuroclr(x)
+        return {"h": h, "z": z}
 
 # --------------------------
 # Your ResNet1D head (verbatim)