move to cuda

ml-energy · Dec 13, 2024 · 4140af5 · 4140af5
1 parent 856b48e
commit 4140af5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/zeus/utils/framework.py b/zeus/utils/framework.py
@@ -122,7 +122,7 @@ def all_reduce(
             return object
 
         # wrap object in a tensor
-        tensor = torch.Tensor(object, device="cuda")
+        tensor = torch.Tensor(object).cuda()
 
         # determine operation
         if operation == "sum":