[Fix] enable weights_only in torch.load to fix tests

Author: GaetanLepage

mmengine/runner/checkpoint.py

@@ -344,7 +344,7 @@ def load_from_local(filename, map_location):
     filename = osp.expanduser(filename)
     if not osp.isfile(filename):
         raise FileNotFoundError(f'{filename} can not be found.')
-    checkpoint = torch.load(filename, map_location=map_location)
+    checkpoint = torch.load(filename, map_location=map_location, weights_only = False)
     return checkpoint
 
 
@@ -412,7 +412,7 @@ def load_from_pavi(filename, map_location=None):
     with TemporaryDirectory() as tmp_dir:
         downloaded_file = osp.join(tmp_dir, model.name)
         model.download(downloaded_file)
-        checkpoint = torch.load(downloaded_file, map_location=map_location)
+        checkpoint = torch.load(downloaded_file, map_location=map_location, weights_only = False)
     return checkpoint
 
 
@@ -435,7 +435,7 @@ def load_from_ceph(filename, map_location=None, backend='petrel'):
     file_backend = get_file_backend(
         filename, backend_args={'backend': backend})
     with io.BytesIO(file_backend.get(filename)) as buffer:
-        checkpoint = torch.load(buffer, map_location=map_location)
+        checkpoint = torch.load(buffer, map_location=map_location, weights_only = False)
     return checkpoint
 
 
@@ -504,7 +504,7 @@ def load_from_openmmlab(filename, map_location=None):
         filename = osp.join(_get_mmengine_home(), model_url)
         if not osp.isfile(filename):
             raise FileNotFoundError(f'{filename} can not be found.')
-        checkpoint = torch.load(filename, map_location=map_location)
+        checkpoint = torch.load(filename, map_location=map_location, weights_only = False)
     return checkpoint
 
 

tests/test_hooks/test_ema_hook.py

@@ -230,7 +230,7 @@ def test_with_runner(self):
         self.assertTrue(
             isinstance(ema_hook.ema_model, ExponentialMovingAverage))
 
-        checkpoint = torch.load(osp.join(self.temp_dir.name, 'epoch_2.pth'))
+        checkpoint = torch.load(osp.join(self.temp_dir.name, 'epoch_2.pth'), weights_only = False)
         self.assertTrue('ema_state_dict' in checkpoint)
         self.assertTrue(checkpoint['ema_state_dict']['steps'] == 8)
 
@@ -245,7 +245,7 @@ def test_with_runner(self):
         runner.test()
 
         # Test load checkpoint without ema_state_dict
-        checkpoint = torch.load(osp.join(self.temp_dir.name, 'epoch_2.pth'))
+        checkpoint = torch.load(osp.join(self.temp_dir.name, 'epoch_2.pth'), weights_only = False)
         checkpoint.pop('ema_state_dict')
         torch.save(checkpoint,
                    osp.join(self.temp_dir.name, 'without_ema_state_dict.pth'))
@@ -274,7 +274,7 @@ def test_with_runner(self):
         runner = self.build_runner(cfg)
         runner.train()
         state_dict = torch.load(
-            osp.join(self.temp_dir.name, 'epoch_4.pth'), map_location='cpu')
+            osp.join(self.temp_dir.name, 'epoch_4.pth'), map_location='cpu', weights_only = False)
         self.assertIn('ema_state_dict', state_dict)
         for k, v in state_dict['state_dict'].items():
             assert_allclose(v, state_dict['ema_state_dict']['module.' + k])
@@ -287,12 +287,12 @@ def test_with_runner(self):
         runner = self.build_runner(cfg)
         runner.train()
         state_dict = torch.load(
-            osp.join(self.temp_dir.name, 'iter_4.pth'), map_location='cpu')
+            osp.join(self.temp_dir.name, 'iter_4.pth'), map_location='cpu', weights_only = False)
         self.assertIn('ema_state_dict', state_dict)
         for k, v in state_dict['state_dict'].items():
             assert_allclose(v, state_dict['ema_state_dict']['module.' + k])
         state_dict = torch.load(
-            osp.join(self.temp_dir.name, 'iter_5.pth'), map_location='cpu')
+            osp.join(self.temp_dir.name, 'iter_5.pth'), map_location='cpu', weights_only = False)
         self.assertIn('ema_state_dict', state_dict)
 
     def _test_swap_parameters(self, func_name, *args, **kwargs):

tests/test_runner/test_runner.py

@@ -2272,7 +2272,7 @@ def test_checkpoint(self):
         self.assertTrue(osp.exists(path))
         self.assertFalse(osp.exists(osp.join(self.temp_dir, 'epoch_4.pth')))
 
-        ckpt = torch.load(path)
+        ckpt = torch.load(path, weights_only = False)
         self.assertEqual(ckpt['meta']['epoch'], 3)
         self.assertEqual(ckpt['meta']['iter'], 12)
         self.assertEqual(ckpt['meta']['experiment_name'],
@@ -2444,7 +2444,7 @@ def test_checkpoint(self):
         self.assertTrue(osp.exists(path))
         self.assertFalse(osp.exists(osp.join(self.temp_dir, 'epoch_13.pth')))
 
-        ckpt = torch.load(path)
+        ckpt = torch.load(path, weights_only = False)
         self.assertEqual(ckpt['meta']['epoch'], 0)
         self.assertEqual(ckpt['meta']['iter'], 12)
         assert isinstance(ckpt['optimizer'], dict)
@@ -2455,7 +2455,7 @@ def test_checkpoint(self):
         self.assertEqual(message_hub.get_info('iter'), 11)
         # 2.1.2 check class attribute _statistic_methods can be saved
         HistoryBuffer._statistics_methods.clear()
-        ckpt = torch.load(path)
+        ckpt = torch.load(path, weights_only = False)
         self.assertIn('min', HistoryBuffer._statistics_methods)
 
         # 2.2 test `load_checkpoint`