Consider enhancing get_train_stats()

[holgerroth]

Is your feature request related to a problem? Please describe. Currently, the Trainer's get_train_stats() function is returning a fixed set of stats

    def get_train_stats(self) -> dict[str, float]:
        return {"total_epochs": self.state.max_epochs, "total_iterations": self.state.epoch_length}

Describe the solution you'd like Make it configurable true input arguments.

Describe alternatives you've considered Alternatively, anyone needing train stats could access trainer.state directly. However, having a general get_train_stats() call might have more utilities.

Additional context Training stats could be useful in contexts such as federated learning.

[drmhrehman]

I agree with @holgerroth

More general and flexible stats collector is needed in order to support a large variety of single and multimodal trainers.

[Nic-Ma]

I will try to enhance the get_train_stats API with more flexible fields.

Thanks.