You could return the intermediate activations in the forward
as seen here:
def forward(self, x):
x1 = self.layer1(x)
x2 = self.layer2(x1)
out = self.layer3(x2)
return out, x1, x2
and then calculate the losses separately:
out, x1, x2 = model(input)
loss = criterion(out, target) + criterion1(x1, target1) + criterion2(x2, target2)
loss.backward()
or alternatively you could use forward hooks to store the intermediate activations and calculate the loss using them. This post gives you an example on how to use forward hooks.