Enlightened by this thread, I found the solution, which is to use the torch.stack() function. So to get the output y, I did
y = torch.stack([torch.stack([nnFunc.log_softmax(nnFunc.relu(self.linear[i][j](x))) for j,m in enumerate(l)],0) for i,l in enumerate(self.linear)],1)
and then use the torch.permute() function:
return y.permute(2,3,1,0)