dropping inference mode for now

app.py

@@ -42,9 +42,11 @@ def load_model_from_hub(preset, device):
         cache_dir="/tmp/",
     )
 
+    print("HF SAVE DIR:", hf_checkpoint)
+
     model = ScoreFlow(scorenet, device=device, **model_params["PatchFlow"])
     model.load_state_dict(load_file(hf_checkpoint), strict=True)
-
+    model = model.eval().requires_grad_(False)
     return model
 
 

msma.py

@@ -81,7 +81,7 @@ class EDMScorer(torch.nn.Module):
 
         self.register_buffer("sigma_steps", t_steps.to(torch.float64))
 
-    @torch.inference_mode()
+    # @torch.inference_mode()
     def forward(
         self,
         x,
@@ -378,7 +378,7 @@ def train_flow(dataset_path, preset, outdir, epochs, **flow_kwargs):
 
     with open(f"{experiment_dir}/config.json", "w") as f: 
         json.dump(model.config, f, sort_keys=True, indent=4)
-        
+
     # totaliters = int(epochs * train_len)
     pbar = tqdm(range(epochs), desc="Train Loss: ? - Val Loss: ?")
     step = 0