In the following code snippet, I have tried to implement an autograd functional layer interior to a simple model architecture. The function takes a single input x and returns 3 outputs. The whole idea is to see how autograd operates in a multi-output setting. The function behaves as an identity function to the 3 outputs in the forward pass.
I train the model for classification with a randomly initialised (input, target) pair for 10 epochs where the loss can be seen decreasing smoothly. However, when I try to visualize the 3 gradients (corresponding to the 3 outputs) during the backward pass, I find only the gradient corresponding to the last returned tensor has a non-zero value and this happens for any n (n>1). All other gradients are zeroed out.
import numpy as np
import torch
import torch.optim as optim
import torch.nn as nn
from torch.autograd import Function

class Grad_MulOuts(Function):    
    @staticmethod
    def forward(ctx, x):
        input_ = x.clone()
        ctx.save_for_backward(input_)
        out1 = input_
        out2 = out1
        out3 = out2
        return out1, out2, out3

    @staticmethod
    def backward(ctx, *grad_outputs):
        input_, = ctx.saved_tensors
        grad_input1 = grad_outputs[0].clone()
        grad_input2 = grad_outputs[1].clone()
        grad_input3 = grad_outputs[2].clone()
        grad_out = grad_input1+grad_input2+grad_input3
        print('Grad 1:{}'.format(grad_input1))
        print('Grad 2:{}'.format(grad_input2)) 
        print('Grad 3:{}'.format(grad_input3))
        return grad_out
def grad_mulouts(x):
    return Grad_MulOuts.apply(x)

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.l1 = nn.Linear(1024, 100)
        self.classifiers = nn.ModuleList([nn.Linear(100, 1) for i in range(3)])
    def forward(self, x):
        x = self.l1(x)
        out={}
        out[0], out[1], out[2] = grad_mulouts(x)
        for i, cl in enumerate(self.classifiers):
            out[i] = cl(out[i])
        return out

X=torch.randn(100, 1024)
Y=torch.FloatTensor(np.random.choice([0, 1], size=(100,3), p=[1./3, 2./3]))
model=Model()
optimizer=optim.SGD(model.parameters(), lr=0.01, nesterov=True, momentum=0.9)
loss_fn=nn.BCELoss()
batch_size=10
for epoch in range(10):
    model.train()
    runningloss=0.0
    for i in range(0,100,batch_size):
        input_ = X[i:i+batch_size]
        target = Y[i:i+batch_size]
        optimizer.zero_grad()
        with torch.set_grad_enabled(True):
            out = model(input_)
            loss = loss_fn(torch.sigmoid(out[0]), target[:,0])+loss_fn(torch.sigmoid(out[1]), target[:,1])+\
                   loss_fn(torch.sigmoid(out[2]), target[:,2])
            runningloss+=loss.item()
            loss.backward()
            optimizer.step()
    print('Loss:{}'.format(runningloss/batch_size))

This is the code that can show the problem I mentioned above. I know with multiple outputs from a single node, gradients would get added up by default. But, if one wishes to combine the gradients in some other way, torch.autograd.Function might be the only option. But a behaviour as this would make that really difficult.
Maybe, I am missing something. Any suggestions would be really helpful. Thanks in advance.