Remove experimental features committed by mistake

deepspeed_config.json

@@ -23,7 +23,7 @@
     "opt_level": "O2"
   },
   "zero_optimization": {
-    "stage": 2
+    "stage": 1
   },
   "activation_checkpointing": {
     "partition_activations": true,

train_openfold.py

@@ -61,9 +61,9 @@ class OpenFoldWrapper(pl.LightningModule):
         # Compute loss
         loss = self.loss(outputs, batch)
 
-        if(torch.isnan(loss) or torch.isinf(loss)):
-            loss = None
-            logging.warning("loss is NaN. Skipping example...")
+        #if(torch.isnan(loss) or torch.isinf(loss)):
+        #    logging.warning("loss is NaN. Skipping example...")
+        #    loss = loss.new_tensor(0., requires_grad=True)
 
         return {"loss": loss}
 
@@ -117,12 +117,12 @@ def main(args):
         sd = {k[len("module."):]:v for k,v in sd.items()}
         model_module.load_state_dict(sd)
         logging.info("Successfully loaded model weights...")
-    data_module = DummyDataLoader("batch.pickle")
-    #data_module = OpenFoldDataModule(
-    #    config=config.data, 
-    #    batch_seed=args.seed,
-    #    **vars(args)
-    #)
+    #data_module = DummyDataLoader("batch.pickle")
+    data_module = OpenFoldDataModule(
+        config=config.data, 
+        batch_seed=args.seed,
+        **vars(args)
+    )
     data_module.prepare_data()
     data_module.setup()