fix test_optim

DeepRank · Oct 20, 2023 · a6209cb · a6209cb
1 parent a094e2f
commit a6209cb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
@@ -483,7 +483,7 @@ def test_optim(self):
                 dataset_test=dataset,
                 pretrained_model=self.save_path)
 
-        assert isinstance(trainer_pretrained.optimizer, optimizer)
+        assert str(type(trainer_pretrained.optimizer)) == "<class 'torch.optim.adamax.Adamax'>"
         assert trainer_pretrained.lr == lr
         assert trainer_pretrained.weight_decay == weight_decay