[FIX] offload_weight() takes from 3 to 4 positional arguments but 5 were given (#29457)

* use require_torch_gpu * enable on XPU * fix
2024-03-06 10:58:42 +08:00
parent 7b01579f73
commit 00bf44270f
2 changed files with 3 additions and 3 deletions
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -765,7 +765,7 @@ class ModelUtilsTest(TestCasePlus):

    @require_accelerate
    @mark.accelerate_tests
-    @require_torch_accelerator
+    @require_torch_gpu
    def test_from_pretrained_disk_offload_task_model(self):
        model = AutoModel.from_pretrained("hf-internal-testing/tiny-random-gpt2")
        device_map = {
@@ -808,7 +808,7 @@ class ModelUtilsTest(TestCasePlus):

    @require_accelerate
    @mark.accelerate_tests
-    @require_torch_accelerator
+    @require_torch_gpu
    def test_from_pretrained_disk_offload_derived_to_base_model(self):
        derived_model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")