themrzmaster / git-re-basin-pytorch

Git Re-Basin: Merging Models modulo Permutation Symmetries in PyTorch
MIT License
72 stars 9 forks source link

Can't get this to run on a gpu #6

Closed ogkalu2 closed 1 year ago

ogkalu2 commented 2 years ago

If I don't specify cpu when loading the models, I get

File review A = w_a @ w_b.T "Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!"

I now know the linear sum assignment message it can't ask run on a GPU but is it possible to get it to use both ?

affableroots commented 1 year ago

I'm no pytorch expert, but you can use x.cpu() to move something to the CPU, and then x.cuda() to move it back

themrzmaster commented 1 year ago

yea, just move everything to gpu, data and model.