Could you explain what setups you’ve tested so far?
Based on the description of the two setups, the changes are:

PyTorch version
CUDA version
cudnn (if you are using the binaries and the linked cudnn release)
driver version

All of these components might change the memory foot print as well as the performance:

PyTorch might have added more CUDA kernels, thus the CUDA context needs more memory and you would have to reduce the batch size
CUDA as well as cudnn might have increases its memory footprint
depending on the used GPU you might be seeing performance regressions through cudnn (e.g. if the heuristic isn’t fine-tuned for the device and is picking a non-optimal kernel)

Based on the changes, I would not expect that the driver is the culprit, but you might have additional testing, so it would be good to hear it.