Changes model dims

brlp_lite.py

@@ -94,8 +94,8 @@ def load_if(checkpoints_path: Optional[str], network: nn.Module) -> nn.Module:
         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'))
+            network.load_state_dict(torch.load(checkpoints_path))
+            #network.load_state_dict(torch.load(checkpoints_path, map_location='cpu'))
     return network