fix: torch 2.6 weights_only load

Signed-off-by: Hao Lin <[email protected]>

Author: andylin-hao

libero/libero/benchmark/__init__.py

@@ -161,7 +161,7 @@ def get_task_init_states(self, i):
             self.tasks[i].problem_folder,
             self.tasks[i].init_states_file,
         )
-        init_states = torch.load(init_states_path)
+        init_states = torch.load(init_states_path, weights_only=False)
         return init_states
 
     def set_task_embs(self, task_embs):

libero/lifelong/evaluate.py

@@ -251,7 +251,7 @@ def main():
         init_states_path = os.path.join(
             cfg.init_states_folder, task.problem_folder, task.init_states_file
         )
-        init_states = torch.load(init_states_path)
+        init_states = torch.load(init_states_path, weights_only=False)
         indices = np.arange(env_num) % init_states.shape[0]
         init_states_ = init_states[indices]
 

libero/lifelong/metric.py

@@ -104,7 +104,7 @@ def evaluate_one_task_success(
         init_states_path = os.path.join(
             cfg.init_states_folder, task.problem_folder, task.init_states_file
         )
-        init_states = torch.load(init_states_path)
+        init_states = torch.load(init_states_path, weights_only=False)
         num_success = 0
         for i in range(eval_loop_num):
             env.reset()

libero/lifelong/utils.py

@@ -56,7 +56,7 @@ def torch_save_model(model, model_path, cfg=None, previous_masks=None):
 
 
 def torch_load_model(model_path, map_location=None):
-    model_dict = torch.load(model_path, map_location=map_location)
+    model_dict = torch.load(model_path, weights_only=False, map_location=map_location)
     cfg = None
     if "cfg" in model_dict:
         cfg = model_dict["cfg"]