Update Exp2_Global_ocean_simulation/train.py

Exp2_Global_ocean_simulation/train.py

@@ -275,7 +275,7 @@ class Trainer():
 
             if self.params.enable_amp:
                 self.gscaler.update()
-            break
+            # break
 
             tr_time += time.time() - tr_start
 
@@ -308,8 +308,8 @@ class Trainer():
         sample_idx = np.random.randint(len(self.valid_data_loader))
         with torch.no_grad():
             for i, data in enumerate(self.valid_data_loader, 0):
-                if i > 1:
-                    break
+                # if i > 1:
+                #     break
                 inp, tar = map(lambda x: x.to(self.device, dtype=torch.float), data)
                 # gen = self.model(inp)
                 num_steps = params.multi_steps_finetune