val_loss fix and stop sampling recycling iterations in validation

openfold/data/data_modules.py

@@ -283,7 +283,8 @@ class OpenFoldDataLoader(torch.utils.data.DataLoader):
             keyed_probs.append(
                 ("use_clamped_fape", [1 - clamp_prob, clamp_prob])
             )
-            if(self.config.supervised.uniform_recycling):
+
+        if(self.stage == "train" and self.config.supervised.uniform_recycling):
             recycling_probs = [
                 1. / (max_iters + 1) 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, _):