knowing that smaller batches run just fine it shows that you just run out of memory!
what is your model? how many parameters does it have?
you can calculate the amount of memory used by your model and just multiply that by the threads .
for example for the input you have :
256x(224x224x3) x 4 / 1024^2 = 147MB

assuming you are using float32. So 32 threads each reading 256 images, will take 4,704 MB. and this is only the input batch, you need to calculate how much memory your model takes and take that into account as well.
if you want you can use torchstat 28 and get such info and then choose the best course of action.
also recall that in training phase the memory usage is 2~3 times more than what you calculate for your number of parameters since you need to take into account the gradients and optimizes states as well.