prevent side effects in tests

Lightning-AI · Jun 8, 2020 · cc4cb6b · cc4cb6b
1 parent a1ddb86
commit cc4cb6b
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/tests/core/test_decorators.py b/tests/core/test_decorators.py
@@ -15,10 +15,13 @@
 def test_auto_move_data(src_device, dest_device):
     """ Test that the decorator moves the data to the device the model is on. """
 
+    class CurrentModel(EvalModelTemplate):
+        pass
+
     # apply the decorator
-    EvalModelTemplate.forward = auto_move_data(EvalModelTemplate.forward)
+    CurrentModel.forward = auto_move_data(CurrentModel.forward)
 
-    model = EvalModelTemplate()
+    model = CurrentModel()
     model = model.to(dest_device)
     model.prepare_data()
     loader = model.train_dataloader()