tests(serialization): Skip using device_map (requires accelerate)

coreweave · Jun 29, 2024 · afb9606 · afb9606
1 parent edf5d0b
commit afb9606
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/tests/test_serialization.py b/tests/test_serialization.py
@@ -1127,9 +1127,8 @@ def _test_exception_decrypts(self):
         # even in the event of an exception
         encryption = EncryptionParams.random()
         with torch.device("cpu"):
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name, device_map="cpu"
-            )
+            model = AutoModelForCausalLM.from_pretrained(model_name)
+        self.assertEqual(model.device.type, "cpu")
 
         model_sd = model.state_dict()
         model_clone = {k: v.detach().clone() for k, v in model_sd.items()}