Hi, I’m working on a model with two different module (A, B). Module A has 2 conv following by two FC layers and Module B is composed of 5 conv. The model is running on two GPUs.
From my understanding, GPU is allocated when we new a model object with initialization. However, the memory increases about 600M (nvidia-smi) when forwarding thru the first conv(parameter: in_dim=256, out_dim=1024, kernel=3, stride=1) and not increasing any memory when passing thru module B in training phase.
Solved with:
torch.backends.cudnn.benchmark = True
Do not know the reason but problem solved