Code fixes

brlp_lite.py

@@ -56,6 +56,9 @@ from monai.transforms.transform import Transform
 from monai import transforms
 from monai.utils import set_determinism
 from monai.data.meta_tensor import MetaTensor
+import torch.serialization
+
+torch.serialization.add_safe_globals([MetaTensor])
 
 from tqdm import tqdm
 import matplotlib.pyplot as plt
@@ -93,9 +96,9 @@ def load_if(checkpoints_path: Optional[str], network: nn.Module) -> nn.Module:
     if checkpoints_path is not None:
         assert os.path.exists(checkpoints_path), 'Invalid path'
         # Using context manager to allow MetaTensor
-        with torch.serialization.safe_globals([MetaTensor]):
-            network.load_state_dict(torch.load(checkpoints_path))
-            #network.load_state_dict(torch.load(checkpoints_path, map_location='cpu'))
+        #with torch.serialization.safe_globals([MetaTensor]):
+        network.load_state_dict(torch.load(checkpoints_path))
+        #network.load_state_dict(torch.load(checkpoints_path, map_location='cpu'))
     return network
 
 
@@ -490,8 +493,8 @@ def train(
 
         # Save the model after each epoch.
         os.makedirs(output_dir, exist_ok=True)
-        torch.save(discriminator.state_dict(), os.path.join(output_dir, f'discriminator-ep-{epoch + 1}.pth'))
-        torch.save(autoencoder.state_dict(), os.path.join(output_dir, f'autoencoder-ep-{epoch + 1}.pth'))
+        torch.save(discriminator.state_dict(), os.path.join(output_dir, f'discriminator-ep-{epoch}.pth'))
+        torch.save(autoencoder.state_dict(), os.path.join(output_dir, f'autoencoder-ep-{epoch}.pth'))
 
     writer.close()
     print("Training completed and models saved.")