High memory consumption + long training time on a new version

Good day,
I am trying to reproduce some results of this model and the memory consumption is crazy (40gb for 256*256, batch size =3), and the computation time is also very slow, although the number of parameters is very low (~300k). My suggestion that maybe on newer version some changes were done so the code from the older version doesn’t work properly on a new pytorch.
Unfortunately I can’t check the code on an older version because I have CUDA 11 and as far as I understand only pytorch 1,7 is compatible with it.
I am wondering is it possible that on a new version the code runs much slower and consumes a lot of memory? If so, what is the possible reason for that and how it can be fixed?
I also do not exclude the fact the this model has complex connections or wasn’t implemented correctly and simply doesn’t work.

The PyTorch conda binaries and pip wheels ship with their own CUDA runtime, so you could install different versions using virtual environments on the same system. Your local CUDA toolkit would be used, if you build PyTorch from source or compile a custom CUDA extension.
Since you suspect that framework changes are responsible for the increased memory usage, this would allow you to compare different release versions.

thank you for the reply. I tried to install 1.4 version via
conda install pytorch==1.4.0 torchvision==0.5.0 cudatoolkit=10.1 -c pytorch
The installation was successful, cuda_is_available returns True, but the computation takes forever even for the models for that it worked really fast on 1.7 version. My GPU model is Nvidia A100. nvidia-smi says that 500mb are busy, before it was about 2-3GB for the prediction. Could it be possible that older pytorch versions aren’t compatible with this gpu?

They are compatible, but the CUDA JIT will compile all kernels for your architecture, since PyTorch 1.4 didn’t ship with sm_80, which is needed for the A100 you are using.

Thank you for your help. I asked the developer and they confirmed that the model requires a lot of memory and time for training, so I do not think that there is an issue with the performance because of another pytorch version