Register TDA target parameters as loss module buffers

luigibonati / mlcolvar

A unified framework for machine learning collective variables for enhanced sampling simulations

MIT License

91 stars 24 forks source link

Closed recisic closed 2 months ago

recisic commented 2 months ago

This PR introduces registration of the TDA target parameters (target_centers and target_sigmas) as buffers within the TDALoss module.

Eliminates the need to convert both parameters to tensors and move them to the current device for each loss calculation step
Enables retrieval of target parameters from saved DeepTDA CV checkpoints or traced modules

EnricoTrizio commented 2 months ago

Ciao Juno, that looks great! I'll merge and close. Thanks for the help 🚀