ptrblck:

Is your training running fine on a single GPU?
If the loss is consuming so much memory, I would assume a single GPU should consume the same amount.


yes on a single GPU with smaller batch-size works fine.