From 34861a259b3f9691931a4f3ae7c57b2e1ae823e3 Mon Sep 17 00:00:00 2001 From: Fanli Lin Date: Tue, 20 Aug 2024 01:05:12 -0700 Subject: [PATCH] enable --- tests/regression/test_regression.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/regression/test_regression.py b/tests/regression/test_regression.py index f4a1910a50..e3189ee65b 100644 --- a/tests/regression/test_regression.py +++ b/tests/regression/test_regression.py @@ -187,7 +187,7 @@ def save_model(model, name, force=False): def load_output(name): filename = os.path.join(REGRESSION_DIR, name, "output.pt") - return torch.load(filename) + return torch.load(filename, map_location=infer_device()) @pytest.mark.regression