Fix test for lightning 2.1

LarsKue · Oct 27, 2023 · 935ffdd · 935ffdd
1 parent b87f433
commit 935ffdd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/module_tests/test_trainable.py b/tests/module_tests/test_trainable.py
@@ -132,7 +132,7 @@ def test_continue_training(dummy_model):
     print("Finished Training. Loading Checkpoint.")
     checkpoint = find_checkpoint()
 
-    trained_model = dummy_model.load_from_checkpoint(checkpoint)
+    trained_model = dummy_model.__class__.load_from_checkpoint(checkpoint)
 
     print("Continuing Training.")
     trained_model.fit(