Thank you for the reply Patrick!
I’m still a bit confused to understand however how my gradient will be calculated w.r.t the observations in the batch with the for-loop.
For example if my nn.Module has ModuleDict(Block1A, Block1B, block1C) -> Block2 -> Block3 -> Out
And in my input-batch I have say (all) observations corresponding to Block1A then wouldn’t the following produce gradients:
One-by-one:
x = torch.cat([self.moduledict['Block1A'](x[i].unsqueeze(0)) for i in range(len(x))])
Together:
x = self.moduledict['Block1A'](x)
Just to illustrate with a code-example:
import torch
torch.manual_seed(0)
import torch.nn as nn
import torch.nn.functional as F

class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.pre_heads = nn.ModuleDict()
        for head_name in ['a','b','c']:
                self.pre_heads[head_name] = nn.Sequential(
                    nn.Conv2d(16,16, kernel_size=1),
                    nn.ReLU()
                )
        self.out = nn.Sequential(
            Flatten(),
            nn.Linear(16*4*4,1)
        )
        
    def forward(self, x, head_name):           
        x = torch.cat(
            [self.pre_heads[j](x[i].unsqueeze(0)) for i, j in enumerate(head_name)],
            dim=0
        )
        x = torch.squeeze(torch.sigmoid(self.out(x)))
        return x

mymodel = Model()
sample_input = torch.ones(4,16,4,4)
output = mymodel(sample_input, head_name=['a', 'a', 'a', 'a'])
loss = F.binary_cross_entropy(output, torch.Tensor([1,1,1,1]))
loss.backward()
mymodel.pre_heads['a'][0].weight.grad.sum()
# tensor(2.3628)

v.s.
class Model(nn.Module):        
    # Only changed Forward
    def forward(self, x):           
        x = self.pre_heads['a'](x)
        x = torch.squeeze(torch.sigmoid(self.out(x)))
        return x

mymodel = Model()
sample_input = torch.ones(4,16,4,4)
output = mymodel(sample_input])
loss = F.binary_cross_entropy(output, torch.Tensor([1,1,1,1]))
loss.backward()
mymodel.pre_heads['a'][0].weight.grad.sum()
# tensor(-0.9457)