update test

Lightning-AI · Jul 3, 2024 · d2589fc · d2589fc
1 parent 137c7b9
commit d2589fc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/tests_pytorch/trainer/test_trainer.py b/tests/tests_pytorch/trainer/test_trainer.py
@@ -1766,7 +1766,7 @@ def current_memory():
     trainer.fit(model)
 
     assert trainer.strategy.model is model
-    assert list(trainer.optimizers[0].state.values())[0]["exp_avg_sq"].device == torch.device("cpu")
+    assert list(trainer.optimizers[0].state.values())[0]["exp_avg_sq"].device == torch.device("cuda", 0)
     assert trainer.callback_metrics["train_loss"].device == torch.device("cpu")
 
     assert current_memory() <= initial