Hello -
My model looks something like this:
class Model(nn.Module):
    def __init__(self,labels):
        super(Model,self).__init__()
        self.labels=labels
        n_labels=len(labels)
        label_models=nn.ModuleList()
        for n_models in range(n_labels):
            label_models.append(MLP())
        self.labelspecific_models=label_models

    def forward(self,data):
        prediction=torch.zeros(batch_size) '''batch_size=10'''
        prediction.require_grad=True
        for index,label in enumerate(self.labels):
            model_inputs=data[label][0]
            contribution_index=data[label][1]
            label_outputs=self.labelspecific_models[index].forward(model_inputs)
            for index,label_output in enumerate(label_outputs):
                prediction[contribution_index]+=label_output
        return prediction

As you can see, my input data is fed into different models and then joined together depending on their corresponding contribution index. My question is, is there any issue here with adding the outputs to ‘prediction’ and returning that value and training accordingly? Will the gradients be computed without issue and train properly. I seem to have issues but I can’t identify if its from this part. Thanks!