Loss does not decrease when the model is called twice

Hello everyone, I have a model, say, net, where its forward function contains an if else statement inside. It is something like this:

def forward(self, x):
   if first case:
       self.module_1() #run this module only
   else:
       self.module_2() #run this module only

The special thing is that I need to call the forward function of net two times. Moreover, the later one uses the first one’s output as the input. It is something like this:

output = net(input for first case)
second_output = net(output from the first call as the input)

While the code is running OK, what I found strange with this is that when I run training the loss does not decrease. This means something is wrong but I don’t know what it is and why.

Any explanation? Thanks.

Actually I found the real reason behind this. In between these two lines:

output = net(input for first case)
second_output = net(output from the first call as the input)

I actually have several extra simple operations over the output tensor in the middle. Strangely enough, with that back-propagation does not work properly. My solution is to wrap these extra operations into a nn.module. With this it works normally again.

That is it. I hope my experience is useful to someone.

2 Likes