Modalities / modalities

A framework for training multimodal foundation models.
MIT License
39 stars 3 forks source link

Improve indexation of the cummulated_loss_tensor in the Trainer class #123

Open le1nux opened 3 months ago

le1nux commented 3 months ago

As seen here, https://github.com/Modalities/modalities/blob/a599005e646a0ec1db343fbe1ef3be7b53cb75f9/src/modalities/trainer.py#L144, the cumulated_loss_and_gradient_norm is indexed directly in the train function, adding unnecessary complexity to the code.

A better solution would be to encapsulate the cumulated_loss_and_gradient_norm within a dedicated class that implements also the reduce operations.