device type

Genentech · Jul 30, 2024 · 1d2e4e3 · 1d2e4e3
1 parent 9c1aecb
commit 1d2e4e3
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/beignet/func/_partition.py b/src/beignet/func/_partition.py
@@ -1180,8 +1180,11 @@ def fn(
                 device=parameter.device,
             )
 
+        positions = positions.to(device=device)
+        unit_size = unit_size.to(device=device)
+
         hashes = torch.sum(
-            (positions / unit_size).to(dtype=torch.int32)
+            (positions / unit_size).to(dtype=torch.int32, device=device)
             * _hash_constants(spatial_dimension, units_per_side).to(
                 device=positions.device
             ),