Comment by kookamamie

> I would guess the constant memory allocation and frees in the training loop are the bottleneck

No, the bottleneck would be not utilizing the idling GPU.