Interesting, is the script doing some kind of runtime batch size adjustment after checking the amount of GPU memory? I would be surprised if the batch size is the same across GPUs with different memory usage unless something like precision of the data types is changed.
Yes, I don’t think there is a user-facing controllable way (or even a PyTorch controllable way short of rebuilding the libraries) to change the remaining 1GiB+ of memory usage as this is dedicated to the library code (native kernels, cuDNN, cuBLAS, etc.) that cannot be removed.