test: move the pretrained model to the same device as current model b…

…efore assertion (#19668)
Lightning-AI · Mar 18, 2024 · f4178d5 · f4178d5
1 parent b3bd871
commit f4178d5
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/tests_pytorch/models/test_restore.py b/tests/tests_pytorch/models/test_restore.py
@@ -446,6 +446,7 @@ def test_load_model_from_checkpoint(tmp_path, model_template):
         "limit_test_batches": 2,
         "callbacks": [ModelCheckpoint(dirpath=tmp_path, monitor="val_loss", save_top_k=-1)],
         "default_root_dir": tmp_path,
+        "accelerator": "cpu",
     }
 
     # fit model