I don’t think there is a reliable way to estimate the memory usage properly.
Especially, if you are using cudnn in the benchmark mode (via torch.backends.cudnn.benchmark = True), as cudnn will profile various algorithms for your use case and will select the fastest, which fits into your device. E.g. a faster algorithm might be picked, if you have enough memory left, and will be skipped, if not enough memory is available.
This could mean that the memory usage won’t scale with the increase in input size as expected.
Of course you could try to calculate a theoretical memory usage, but I would just run some tests using different shapes to get the estimated memory usage.