Moving solver between devices (multiple GPUs)

rgl-epfl / cholespy

An easily integrable Cholesky solver on CPU and GPU

BSD 3-Clause "New" or "Revised" License

210 stars 18 forks source link

Moving solver between devices (multiple GPUs) #32

Open ArmanMaesumi opened 5 months ago

ArmanMaesumi commented 5 months ago

Hi, thank you for the great package.

I'm wondering if it's possible to move a CholeskySolverF between different devices. Currently the solver is placed on the device corresponding to its input data.

In my context I have many cholesky solver objects in a large dataset, and I need to readily move them to different gpus inside a torch training loop. Is this possible?

Thanks.