Fix missing initialization of contrastive loss

psaegert · Jan 29, 2025 · 9d133e5 · 9d133e5
1 parent 154023d
commit 9d133e5
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/flash_ansr/train/train.py b/src/flash_ansr/train/train.py
@@ -398,6 +398,8 @@ def _train_batch(self, batch: dict[str, torch.Tensor], numeric_prediction_loss_w
             # Calculate the loss
             ce_loss: torch.Tensor = self.cross_entropy_loss(flat_logits, flat_labels)
 
+            contrastive_loss = torch.tensor(0, device=self.model.device, dtype=torch.float32)
+
             if contrastive_loss_weight > 0:
                 # Use memory (embeddings of encoder) that the model cached during the forward pass
                 contrastive_loss = self.contrastive_loss_fn(self.model.memory.reshape(self.model.memory.shape[0], -1), skeleton_hashes)
@@ -531,6 +533,8 @@ def _validate(self, val_dataset: FlashANSRDataset, numeric_prediction_loss_weigh
                 # Calculate the loss
                 ce_loss: torch.Tensor = self.cross_entropy_loss(flat_logits, flat_labels)
 
+                contrastive_loss = torch.tensor(0, device=self.model.device, dtype=torch.float32)
+
                 if contrastive_loss_weight > 0:
                     # Use memory (embeddings of encoder) that the model cached during the forward pass
                     contrastive_loss = self.contrastive_loss_fn(self.model.memory.reshape(self.model.memory.shape[0], -1), skeleton_hashes)