Remove hpu:X notation untill fully supported by bridge (huggingface#637)

jychen21 · Feb 27, 2024 · 2f9ebae · 2f9ebae
1 parent 1996e05
commit 2f9ebae
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/optimum/habana/accelerate/state.py b/optimum/habana/accelerate/state.py
@@ -76,7 +76,8 @@ def __init__(self, cpu: bool = False, **kwargs):
                 self.process_index = rank
                 self.local_process_index = local_rank
                 if self.device is None:
-                    self.device = torch.device("hpu", self.local_process_index)
+                    # TODO: replace by `torch.device("hpu", self.local_process_index)` when hpu:x is supported
+                    self.device = torch.device("hpu")
             else:
                 self.distributed_type = GaudiDistributedType.NO
                 self.num_processes = 1