Added validation step

JVGD · Oct 23, 2021 · eac8272 · eac8272
1 parent 1af1395
commit eac8272
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/experiment1/model.py b/experiment1/model.py
@@ -43,3 +43,17 @@ def training_step(self, batch, batch_idx, *args, **kwargs) -> T.Tensor:
         self.log("loss/train", loss, prog_bar=True, on_step=False, on_epoch=True)
         return loss
 
+    def validation_step(self, batch, batch_idx, *args, **kwargs) -> T.Tensor:
+        # Unpacking
+        samples = batch["samples"]
+        targets = batch["targets"]
+
+        # Forward
+        targets_pred = self(samples)
+
+        # Loss
+        loss = self.criteria(targets, targets_pred)
+
+        # Logging
+        self.log("loss/valid", loss, prog_bar=True, on_step=False, on_epoch=True)
+        return loss