I checked it out and there was a problem in autograd indeed. We haven’t been freeing some buffers soon enough. I’ve opened a PR 361 with a fix. After that commit batch_size 256 uses the same amount of memory (4,7GB on my GPU) as batch_size 16 before.
Thanks for posting the repro!