npu mlu xpu

Oneflow-Inc · Aug 21, 2024 · 1aaccff · 1aaccff
1 parent 0fb057d
commit 1aaccff
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/libai/utils/distributed.py b/libai/utils/distributed.py
@@ -228,7 +228,13 @@ def device_type(self):
         return self._device_type
 
     def set_device_type(self, device_type):
-        assert device_type in ["cpu", "cuda"], f"not supported for {device_type}"
+        assert device_type in [
+            "cpu",
+            "cuda",
+            "npu",
+            "xpu",
+            "mlu",
+        ], f"not supported for {device_type}"
         self._device_type = device_type
 
     def get_layer_ranks(self, layer_idx):