torch2.4.0 compat

Saurav Agarwal · Saurav Agarwal · commit 667908f91231 · 2024-08-19T22:53:40.000-04:00
diff --git a/python/coverage_control/algorithms/controllers.py b/python/coverage_control/algorithms/controllers.py
@@ -125,7 +125,7 @@ def __init__(self, config: dict, params: Parameters, env: CoverageSystem):
         self.actions_std = self.model.actions_std.to(self.device)
         self.model = self.model.to(self.device)
         self.model.eval()
-        # self.model = torch.compile(self.model, dynamic=True)
+        self.model = torch.compile(self.model, dynamic=True)
 
     def step(self, env):
         """
diff --git a/python/coverage_control/io_utils.py b/python/coverage_control/io_utils.py
@@ -76,7 +76,7 @@ def load_tensor(path: str) -> torch.tensor:
         if not os.path.exists(path):
             raise FileNotFoundError(f"IOUtils::load_tensor: File not found: {path}")
         # Load data
-        data = torch.load(path)
+        data = torch.load(path, weights_only=True)
         # Extract tensor if data is in jit script format
 
         if isinstance(data, torch.jit.ScriptModule):
diff --git a/python/coverage_control/nn/models/lpac.py b/python/coverage_control/nn/models/lpac.py
@@ -82,32 +82,32 @@ def forward(self, data: torch_geometric.data.Data) -> torch.Tensor:
 
     def load_compiled_state_dict(self, model_state_dict_path: str) -> None:
         # remove _orig_mod from the state dict keys
-        state_dict = torch.load(model_state_dict_path)
+        state_dict = torch.load(model_state_dict_path, weights_only=True)
         new_state_dict = {}
         for key in state_dict.keys():
             new_state_dict[key.replace("_orig_mod.", "")] = state_dict[key]
-        self.load_state_dict(new_state_dict, strict=True)
+        self.load_state_dict(new_state_dict, strict=True, weights_only=True)
 
     def load_model(self, model_state_dict_path: str) -> None:
         """
         Load the model from the state dict
         """
-        self.load_state_dict(torch.load(model_state_dict_path), strict=True)
+        self.load_state_dict(torch.load(model_state_dict_path), strict=True, weights_only=True)
 
     def load_model_state_dict(self, model_state_dict: dict) -> None:
         """
         Load the model from the state dict
         """
-        self.load_state_dict(model_state_dict, strict=True)
+        self.load_state_dict(model_state_dict, strict=True, weights_only=True)
 
     def load_cnn_backbone(self, model_path: str) -> None:
         """
         Load the CNN backbone from the model path
         """
-        self.load_state_dict(torch.load(model_path).state_dict(), strict=True)
+        self.load_state_dict(torch.load(model_path).state_dict(), strict=True, weights_only=True)
 
     def load_gnn_backbone(self, model_path: str) -> None:
         """
         Load the GNN backbone from the model path
         """
-        self.load_state_dict(torch.load(model_path).state_dict(), strict=True)
+        self.load_state_dict(torch.load(model_path).state_dict(), strict=True, weights_only=True)
