val_loss fix and stop sampling recycling iterations in validation

openfold/data/data_modules.py

@@ -283,13 +283,14 @@ class OpenFoldDataLoader(torch.utils.data.DataLoader):
             keyed_probs.append(
                 ("use_clamped_fape", [1 - clamp_prob, clamp_prob])
             )
-            if(self.config.supervised.uniform_recycling):
-                recycling_probs = [
-                    1. / (max_iters + 1) for _ in range(max_iters + 1)
-                ]
-                keyed_probs.append(
-                    ("no_recycling_iters", recycling_probs)
-                )
+
+        if(self.stage == "train" and self.config.supervised.uniform_recycling):
+            recycling_probs = [
+                1. / (max_iters + 1) for _ in range(max_iters + 1)
+            ]
+            keyed_probs.append(
+                ("no_recycling_iters", recycling_probs)
+            )
         else:
             recycling_probs = [
                 0. for _ in range(max_iters + 1)

train_openfold.py

@@ -66,7 +66,7 @@ class OpenFoldWrapper(pl.LightningModule):
 
         # Compute loss
         loss = self.loss(outputs, batch)
-
+        self.log("loss", loss)
         return {"loss": loss}
 
     def validation_step(self, batch, batch_idx):
@@ -79,6 +79,7 @@ class OpenFoldWrapper(pl.LightningModule):
         outputs = self(batch)
         batch = tensor_tree_map(lambda t: t[..., -1], batch)
         loss = self.loss(outputs, batch)
+        self.log("val_loss", loss, prog_bar=True)
         return {"val_loss": loss}
 
     def validation_epoch_end(self, _):