test: move the pretrained model to the same device as current model before assertion (#19668)

2024-03-18 23:23:02 +05:30 · 2024-03-18 23:23:02 +05:30 · f4178d5d5a
parent b3bd871709
commit f4178d5d5a
1 changed files with 1 additions and 0 deletions
--- 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