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
130 stars 0 forks source link

Add cuda tests #104

Closed ValerianRey closed 3 weeks ago

ValerianRey commented 3 weeks ago

This PR adds the possibility to run unit tests on cuda (fixes #97).

As a result, a test that was already slow became even slower on cuda, so this PR also changes this test:

Lastly, while developing this PR, I found several small issues (sometimes unrelated to the device):

In my opinion the only commit worth a changelog entry is fd1e501a48ecb3e8d5f6d1fa1f55f1e68ff99bbd, so for me the changelog is ready.