Oh, sorry for not communicating this clearly earlier – but I think the issue that I raised in this thread is that PyTorch could not reserve as much memory as what it did for the one without “–shm_cuda” arg.
If you take a look at my sample output before throwing the error:
[P1] [GPU Mem (MB)] total: 16130.5, reserved = 9858.0, allocated: 9856.0625, free: 1.9375, step: 68
...
RuntimeError: CUDA out of memory. Tried to allocate 64.00 MiB (GPU 0; 15.75 GiB total capacity; 9.63 GiB already allocated; 59.88 MiB free; 9.63 GiB reserved in total by PyTorch)

After reserving for ~10GB, it went out of memory even though the total available is ~16GB (there was no other process using the GPU). Meanwhile, in the earlier case you could see that the program managed to finish with reserving ~13GB
[P0] [GPU Mem (MB)] total: 16130.5, reserved = 13826.0, allocated: 13824.0625, free: 1.9375, step: 99

The memory usages are indeed the same for both use cases. It’s just that somehow PyTorch couldn’t reserve the same amount of memory.
I purposely made the sample code to get out of memory. Hence if you have GPUs with diff total memory, prob need to reduce/increase the number of steps to be able to reproduce.