fixed LoadTorchMLP CPU issue in torch_mlp.py (#12)

* fixed LoadTorchMLP CPU issue in torch_mlp.py

* blackify

* updated run_tests.yml

* blackify w/ 24.x

* linting error fixed

lanfactory/trainers/torch_mlp.py

@@ -704,7 +704,13 @@ def __init__(self, model_file_path=None, network_config=None, input_dim=None):
             input_shape=self.input_dim,
             generative_model_id=None,
         )
-        self.net.load_state_dict(torch.load(self.model_file_path))
+        # self.net.load_state_dict(torch.load(self.model_file_path))
+        if not torch.cuda.is_available():
+            self.net.load_state_dict(
+                torch.load(self.model_file_path, map_location=torch.device("cpu"))
+            )
+        else:
+            self.net.load_state_dict(torch.load(self.model_file_path))
         self.net.to(self.dev)
         self.net.eval()
 
@@ -768,7 +774,13 @@ def __init__(self, model_file_path=None, network_config=None, input_dim=None):
             input_shape=self.input_dim,
             generative_model_id=None,
         )
-        self.net.load_state_dict(torch.load(self.model_file_path))
+        # self.net.load_state_dict(torch.load(self.model_file_path))
+        if not torch.cuda.is_available():
+            self.net.load_state_dict(
+                torch.load(self.model_file_path, map_location=torch.device("cpu"))
+            )
+        else:
+            self.net.load_state_dict(torch.load(self.model_file_path))
         self.net.to(self.dev)
 
     # AF-TODO: Seemingly LoadTorchMLPInfer is still not callable !