I\'m currently working on porting code from Keras to PyTorch. I\'m working with many GPUs and CPUs so it\'s important to have batch generation happening in parallel. I\'m