TST Enable regression tests on XPU (#2019)

huggingface · Aug 20, 2024 · 25ab6c9 · 25ab6c9
1 parent b4cf1b3
commit 25ab6c9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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