Update test_finetune

LukasHedegaard · Apr 17, 2023 · c4a7a42 · c4a7a42
1 parent dfdc2e8
commit c4a7a42
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_finetune.py b/tests/test_finetune.py
@@ -9,7 +9,7 @@
 
 from ride.core import RideModule
 from ride.finetune import Finetunable
-from ride.optimizers import AdamWCyclicLrOptimizer
+from ride.optimizers import SgdReduceLrOnPlateauOptimizer
 from ride.utils.checkpoints import get_latest_checkpoint
 from ride.utils.utils import AttributeDict
 
@@ -20,7 +20,7 @@
 
 
 class DummyModule(
-    RideModule, Finetunable, DummyRegressionDataLoader, AdamWCyclicLrOptimizer
+    RideModule, Finetunable, DummyRegressionDataLoader, SgdReduceLrOnPlateauOptimizer
 ):
     def __init__(self, hparams):
         self.l1 = torch.nn.Linear(self.input_shape[0], self.hparams.hidden_dim)