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

awaelchli · fnhirwa · lantiga · commit f4cefa07ecbb · 2024-04-11T09:09:21.000-04:00
Co-authored-by: Felix Hirwa Nshuti &lt;hirwanshutiflx@gmail.com&gt;
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(tmpdir, model_template):
         "limit_test_batches": 2,
         "callbacks": [ModelCheckpoint(dirpath=tmpdir, monitor="val_loss", save_top_k=-1)],
         "default_root_dir": tmpdir,
+        "accelerator": "cpu",
     }
 
     # fit model

Original file line number	Diff line number	Diff line change
`@@ -446,6 +446,7 @@ def test_load_model_from_checkpoint(tmpdir, model_template):`
`446`	`446`	`"limit_test_batches": 2,`
`447`	`447`	`"callbacks": [ModelCheckpoint(dirpath=tmpdir, monitor="val_loss", save_top_k=-1)],`
`448`	`448`	`"default_root_dir": tmpdir,`
	`449`	`+ "accelerator": "cpu",`
`449`	`450`	`}`
`450`	`451`
`451`	`452`	`# fit model`