new ckpt

.gitignore

@@ -1,6 +1,11 @@
 # Mac
 .DS_Store
 
+# PyCharm
 .idea/
+
+# Jupyter notebooks
 .ipynb_checkpoints/
+
+# Python
 __pycache__/

factories.py

@@ -212,29 +212,24 @@ class EvalModel(torch.nn.Module):
     def __init__(self, model_name: str, ckpt_pth: str = "", device_str: str = "cpu") -> None:
         """Load the model we want to evaluate."""
         super().__init__()
-        self.base_name = model_name
+        self.name = model_name
         self.ckpt_pth = ckpt_pth
-        self.name = self.base_name
-        if self.base_name not in self.all_models:
-            raise ValueError(f"{self.base_name} is unavailable.")
-        if self.base_name == "unext_emb_physics_config_C":
+        if self.name not in self.all_models:
+            raise ValueError(f"{self.name} is unavailable.")
+        if self.name == "unext_emb_physics_config_C":
             if self.ckpt_pth == "":
-                self.ckpt_pth = "ckpt/ram_ckp_10.pth.tar"
-            self.model = get_model(model_name=self.base_name,
+                self.ckpt_pth = "ckpt/ram.pth.tar"
+            self.model = get_model(model_name=self.name,
                                    device='cpu',
                                    **DEFAULT_MODEL_PARAMS)
 
-            # load model checkpoint
-            state_dict = torch.load(self.ckpt_pth, map_location=lambda storage, loc: storage)[
-                'state_dict']  # load on cpu
+            # load model checkpoint on cpu
+            state_dict = torch.load(self.ckpt_pth, map_location=lambda storage, loc: storage)
+
             self.model.load_state_dict(state_dict)
             self.model.to(device_str)
             self.model.eval()
 
-            # add epoch in the model name
-            epoch = torch.load(self.ckpt_pth, map_location=lambda storage, loc: storage)['epoch']
-            self.name = self.name + f"+{epoch}"
-
     def forward(self, y: torch.Tensor, physics: torch.nn.Module) -> torch.Tensor:
         return self.model(y, physics=physics)
 
@@ -250,9 +245,8 @@ class BaselineModel(torch.nn.Module):
 
     def __init__(self, model_name: str, device_str: str = "cpu") -> None:
         super().__init__()
-        self.base_name = model_name
+        self.name = model_name
         self.ckpt_pth = ""
-        self.name = self.base_name
         if self.name not in self.all_baselines:
             raise ValueError(f"{self.name} is unavailable.")
         elif self.name == "DPIR":