Fix failing tests

Lightning-AI · Mar 23, 2021 · 210fd87 · 210fd87
1 parent 23b21c5
commit 210fd87
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/plugins/test_double_plugin.py b/tests/plugins/test_double_plugin.py
@@ -125,5 +125,6 @@ def test_double_precision(tmpdir, boring_model):
     trainer.fit(model)
     trainer.test(model)
     trainer.predict(model)
+    torch.set_grad_enabled(True)  # trainer.predict kills gradient
 
     assert model.training_step == original_training_step