Skip to content

Commit c2d9739

Browse files
HAOCHENYEMGAMZ
authored andcommitted
[Test] Fix ut of runner
1 parent 74f989b commit c2d9739

File tree

2 files changed

+4
-3
lines changed

2 files changed

+4
-3
lines changed

tests/test_runner/test_checkpoint.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -356,6 +356,7 @@ def load_from_abc(filename, map_location):
356356
assert loader.__name__ == 'load_from_abc'
357357

358358

359+
@patch.dict(sys.modules, {'petrel_client': MagicMock()})
359360
def test_save_checkpoint(tmp_path):
360361
model = Model()
361362
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)

tests/test_runner/test_runner.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2271,7 +2271,7 @@ def test_checkpoint(self):
22712271
self.assertTrue(osp.exists(path))
22722272
self.assertFalse(osp.exists(osp.join(self.temp_dir, 'epoch_4.pth')))
22732273

2274-
ckpt = torch.load(path)
2274+
ckpt = torch.load(path, weights_only=False)
22752275
self.assertEqual(ckpt['meta']['epoch'], 3)
22762276
self.assertEqual(ckpt['meta']['iter'], 12)
22772277
self.assertEqual(ckpt['meta']['experiment_name'],
@@ -2445,7 +2445,7 @@ def test_checkpoint(self):
24452445
self.assertTrue(osp.exists(path))
24462446
self.assertFalse(osp.exists(osp.join(self.temp_dir, 'epoch_13.pth')))
24472447

2448-
ckpt = torch.load(path)
2448+
ckpt = torch.load(path, weights_only=False)
24492449
self.assertEqual(ckpt['meta']['epoch'], 0)
24502450
self.assertEqual(ckpt['meta']['iter'], 12)
24512451
assert isinstance(ckpt['optimizer'], dict)
@@ -2456,7 +2456,7 @@ def test_checkpoint(self):
24562456
self.assertEqual(message_hub.get_info('iter'), 11)
24572457
# 2.1.2 check class attribute _statistic_methods can be saved
24582458
HistoryBuffer._statistics_methods.clear()
2459-
ckpt = torch.load(path)
2459+
ckpt = torch.load(path, weights_only=False)
24602460
self.assertIn('min', HistoryBuffer._statistics_methods)
24612461

24622462
# 2.2 test `load_checkpoint`

0 commit comments

Comments
 (0)