Training in progress, epoch 1

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d736705fafe165d048f591207c7e984739e7412080ec89b015a457b5084baba6
+oid sha256:a8d0b025faa8c33ff9e377a0ecb3a8066849218e66d79dabdb9c9e1289d3a014
 size 274752173

training_args.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cc65dbe5e03af7e214c614c19b0fda5a22554b261cb74cde5f58af6a97861002
+oid sha256:b65d405be15213e29d93bcec6b150b9a9f69c52594a891222bfb4fa877f3ab74
 size 4027

utils.py

@@ -57,11 +57,14 @@ class ConsistentSentenceRegressor(DistilBertForSequenceClassification):
         # Replace the classifier with a single-neuron linear layer for regression
         self.classifier = torch.nn.Linear(config.dim, config.num_labels)
 
+        self.loss_fn = torch.nn.MSELoss()
+        
         if not freeze_bert:
             return
 
         for param in self.distilbert.parameters():
             param.requires_grad = False
+                
 
     def forward(
         self,
@@ -86,9 +89,13 @@ class ConsistentSentenceRegressor(DistilBertForSequenceClassification):
         )
 
         logits = outputs.logits.squeeze(-1)  # Remove the last dimension to match target tensor shape
-        
+        outputs.logits = logits
+        if labels is not None:
+            # Compute custom loss
+            loss = self.loss_fn(logits, labels)
+            outputs.loss = loss
 
-        return logits
+        return outputs
 
 
 # Set up evaluation metridef get_metrics():