Solved by promach in post #2 
              
              
                Solved by using self.cells = nn.ModuleList([Cell(stride) for i in range(NUM_OF_CELLS)]) instead of just self.cells = [Cell(stride) for i in range(NUM_OF_CELLS)]