search

TorchJD / torchjd

Library for Jacobian descent with PyTorch. It enables optimization of neural networks with multiple losses (e.g. multi-task learning).
https://torchjd.org
MIT License
151 stars 0 forks source link

Run tests on CUDA #97

Closed ValerianRey closed 2 months ago

ValerianRey commented 3 months ago

Currently, all tests are run on cpu. If we mess up with the device field of the tensors, something could break silently. I think it's quite a lot of work, but it would be nice to have tests run on cuda too.