Fix GPU clustering

[jakobnissen]

Commit #6554009 messed up clustering on the GPU, which requires that for all array ops, both arrays are on the same device. Fix this.

NOTE: Marking as draft because currently, the .cpu() calls before torch.histogram destroys performance, so this will have to be fixed.

[jakobnissen]

Okay, this now runs as expected.

pytorch/pytorch#69519 still necessitates that we copy lots of data to and from the GPU which destroys most of the performance from using GPU. I'm quite annoyed by this but I don't see how to get around it. In older version of Vamb, we used the torch.histc function, which does run on GPU, but this can't use weights.

GPU kernels for histogram are available in https://github.com/francois-rozet/torchist, but I'm not sure we should add this dependency.

When running on a dataset with ~120k contigs, the timings are ~660s for CPU and ~812s for GPU, so here GPU is quite a bit slower.

I'm going to merge this now. I'm interested in how performance is for very large datasets (millions of contigs)