I am working a model and it contains multiple instances of a sub network. I am having trouble making each instance of the sub-module learn independently. The gradients are averaged during back prop, making the resulting gradient useless. Below is simplified version of my model:
net_a = ModuleA() 
net_b = ModuleB()

for batch_x, label in dataset: # iterate through dataset
    
    out_b = torch.zeros(10)
    output_arr = []
    for sub_sample in some_sampler(batch): # sampling data along feature dimensions
        out_a = net_a(sub_sample) # Gradient should update `net_a` per slice
        out_b = net_b(out_a, out_b)
        product = out_a.bmm(out_b)
        output_arr.append(product)
        
    output = torch.cat(output_arr)
    loss = F.binary_cross_entropy(output, label, reduction='mean')
    loss.backward()

This model is a bit similar to RNN but I cannot unfold it like a RNN. When back-propagating through the graph, net_a should be receiving a gradient for each sub sample. I believe it has to do with model instances sharing parameters, but I cannot quite grasp how to separate the gradients?