I need to aggregate the outputs from different networks and then do some calculations based on this aggregated outputs, and finally update individual networks. I am unclear about how to keep the outputs which are obtained in the first for loop (from individual networks) and use it in the second for loop to calculate the loss and update individual networks. I very much appreciate if you give me any suggestion.
outputs = []
for i in range(noOfModels):
    optimizer[i].zero_grad()
    tmp = net[i](img) # this will return an output of size 10x10 
    outputs.append(tmp)

outputs = torch.stack(outputs, 2)
x = do_somthing(outputs) # do some calculations with outputs

for i in range(noOfModels):
    loss = criterion(outputs[:,:,i], x)
    loss.backward()
    optimizer[i].step()