Skip to content

Commit

Permalink
Fix GPU clustering
Browse files Browse the repository at this point in the history
A variable which did not support GPU operations was errantly on GPU when
cuda=True. Move to CPU.
  • Loading branch information
jakobnissen committed Nov 24, 2022
1 parent e4c58bf commit dd0406a
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion vamb/cluster.py
Original file line number Diff line number Diff line change
Expand Up @@ -349,7 +349,7 @@ def _smaller_indices(tensor, kept_mask, threshold, cuda):

# If it's on GPU, we remove the already clustered points at this step.
if cuda:
return _torch.nonzero((tensor <= threshold) & kept_mask).flatten()
return _torch.nonzero((tensor <= threshold) & kept_mask).flatten().cpu()
else:
arr = tensor.numpy()
indices = (arr <= threshold).nonzero()[0]
Expand Down

0 comments on commit dd0406a

Please sign in to comment.