Better device_type

PythonOT · Dec 8, 2021 · d02f71a · d02f71a
1 parent c507d3b
commit d02f71a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ot/backend.py b/ot/backend.py
@@ -1688,7 +1688,7 @@ def bitsize(self, type_as):
         return torch.finfo(type_as.dtype).bits
 
     def device_type(self, type_as):
-        return "CPU" if "cpu" in str(type_as.device) else "GPU"
+        return type_as.device.type.replace("cuda", "gpu").upper()
 
     def _bench(self, callable, *args, n_runs=1, warmup_runs=1):
         results = dict()
@@ -2337,7 +2337,7 @@ def bitsize(self, type_as):
         return type_as.dtype.size * 8
 
     def device_type(self, type_as):
-        return "CPU" if "CPU" in type_as.device else "GPU"
+        return self.dtype_device(type_as)[1].split(":")[0]
 
     def _bench(self, callable, *args, n_runs=1, warmup_runs=1):
         results = dict()