device type

Genentech · Jul 30, 2024 · 9c1aecb · 9c1aecb
1 parent 13124ac
commit 9c1aecb
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/beignet/func/_partition.py b/src/beignet/func/_partition.py
@@ -1051,7 +1051,10 @@ def _particles_per_cell(
 
     size, unit_size, per_side, n = _cell_dimensions(dim, size, minimum_size)
 
-    hash_multipliers = _hash_constants(dim, per_side)
+    hash_multipliers = _hash_constants(dim, per_side).to(device=device)
+
+    positions = positions.to(device=device)
+    unit_size = unit_size.to(device=device)
 
     particle_index = torch.tensor(positions / unit_size, dtype=torch.int32, device=device)