From 9c1aecb825ffba960d9cf9c7afbc48b256ba5362 Mon Sep 17 00:00:00 2001 From: Henry Isaacson Date: Mon, 29 Jul 2024 20:06:41 -0400 Subject: [PATCH] device type --- src/beignet/func/_partition.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/beignet/func/_partition.py b/src/beignet/func/_partition.py index fedaed0f17..1a9bfe7747 100644 --- 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)