[Feat] working sinkhorn version and gradient checkpointing in batch dimension

Degiacomi-Lab / molearn

protein conformational spaces meet machine learning

https://degiacomi.org/software/molearn/

GNU General Public License v3.0

40 stars 11 forks source link

Open SCMusson opened 10 months ago

SCMusson commented 10 months ago

A functioning sinkhorn_trainer class Gradient checkpointing module in molearn/utils.py

Two example scripts that hopefully demonstrate both of the above techniques

SCMusson commented 10 months ago

@degiacom This breaks documentation. Do you want to fix this or should I?