provide batch size to self.log() of lightning module to prevent warning

atomistic-machine-learning / schnetpack

SchNetPack - Deep Neural Networks for Atomistic Systems

Other

774 stars 214 forks source link

provide batch size to self.log() of lightning module to prevent warning #598

Closed Maltimore closed 8 months ago

Maltimore commented 8 months ago

this gets rid of annoying warnings of the form

/usr/local/lib/python3.11/site-packages/pytorch_lightning/utilities/data.py:77:
Trying to infer the `batch_size` from an ambiguous collection.
The batch size we found is 100. To avoid any miscalculations,
use `self.log(..., batch_size=batch_size)`.

I assume/hope batch['_idx'] is always available?