Memory increases during training epoch, limiting training set size.

[nkemnitz]

Memory keeps increasing during a single training epoch, after which it drops back to normal levels.
Not sure if bug or intented behavior, but it meant I had to split my training set into smaller chunks and rotate them in and out every "epoch".

I removed all the wandb logging code (especially the per-epoch logging flags) from my regime, but no luck with that.