A simple example is given in below,
class my_model(nn.Module):
     def __init__(self):
         super(my_model,self).__init__()
         self.all_layers = nn.ModuleList()
         for i in range(10):
             layers = []
             layers.append(nn.Linear(10,10))
             layers.append(nn.BatchNorm1d(10))
             self.main = nn.Sequential(*layers)
             self.all_layers.append(self.main)  
      def forward(self,zn,x):
          output = []
          for i in range(10):
              mask = x==i
              temp = zn[mask]
              print(temp.size())
              output.append(self.all_layers[i](temp))
            
          return output



model = my_model()
print(model)
target = torch.randn(60,10)
loss =nn.MSELoss()
input = torch.randn(60,10)
labels = torch.randint(0,10,(60,))
zn = torch.randn(60,10)
predict = model(zn, labels)
print(len(predict))
print(predict)
criteria = loss(predict,target) 

TypeError: expected Tensor as element 0 in argument 0, but got list