AladdinPerzon:

This leads to it being quite unstable as I’m getting OOM error messages sometimes.


Are you getting the OOMs in the same iterations with the large and small batch sizes?
If so, are you changing the input shapes or are you using static shapes?
If you lower the batch size e.g. to 1, are you still seeing OOMs?
Did you check the GPU memory usage via nvidia-smi or with PyTorch methods (e.g. torch.cuda.memory_summary())?
Note that the former would return all allocated memory, i.e. the CUDA context, the used and allocated memory as well as the cache in PyTorch, and the memory usage of all other applications.