Conv module in forward can't alloc

Pytorch is not quite clear about this part, in fact, in this previously asked question (by me):

In come cases (my implementation 2), even if the tensors are dereferenced, their memory are still not freed immediately after dereferencing, for more deep down implementation information, maybe you could consider asking the pytorch development member @ptrblck and @albanD.

I have not tried the jemalloc solution yet, but it is easy enough to install it by simply making it like:

wget https://github.com/jemalloc/jemalloc/releases/download/5.0.1/jemalloc-5.0.1.tar.bz2
tar -jxvf jemalloc-5.0.1.tar.bz2
cd jemalloc-5.0.1
sudo apt-get install autogen autoconf

./autogen.sh
make -j2
sudo make install
sudo ldconfig
cd ../
rm -rf jemalloc-5.0.1 jemalloc-5.0.1.tar.bz2

then export LD_PRELOAD environment variable before running your program:

LD_PRELOAD=/home/mingfeim/packages/jemalloc-5.2.0/lib/libjemalloc.so ./your_script.sh

Give it a try?

Reference: