Normalize Data on GPU

[sadamov]

Motivation
Data normalization can be done on the fly on GPU for each batch. It's faster on GPU than CPU and cleans up the dataset init method.

Implementation
Could very nicely use https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#on-after-batch-transfer to normalize once data is on GPU. Makes sure that you never forget about it (all batches on GPU are normalized).

The stats could be provided by a yaml_object handler that can be accessed on the model's init

[leifdenby]

sounds cool @sadamov, are you thinking this for v0.3.0 or a later release? :)

[sadamov]

This feature is ready in #39 I don't have a strong opinion about the version it should be published in. :)