Update base_trainer.py for handling total_batch_size>1

erfanzar / EasyDeL

Accelerate your training with this open-source library. Optimize performance with streamlined training and serving options with JAX. 🚀

https://easydel.readthedocs.io/en/latest/

Apache License 2.0

167 stars 19 forks source link

Update base_trainer.py for handling total_batch_size>1 #159

Closed s-smits closed 1 month ago

s-smits commented 1 month ago

Currently the expected length of the training run is determined by the number of epochs times the length of the dataset. This modification makes sure it's divided by the total_batch_size to make sure the expected length is matched.