Sudden surge in CPU RAM usage after upgrading pytorch v1.6 to v1.9

Memory Format
3

I will try to upgrade to the latest torch. But the docker image size will also increase significantly. Any pointers to address that?