However, outputs1 dimension depends on combined_feature_map which depends on CIFAR-10 dataset.  So, I cannot simply make the tensor outputs1 to have the shape [batch_size, nb_classes=10].
Please advise.
NUM_OF_IMAGE_CHANNELS = 3

class Cell(nn.Module):
        self.output = torch.zeros([BATCH_SIZE, NUM_OF_IMAGE_CHANNELS, IMAGE_HEIGHT, IMAGE_WIDTH],
                                                requires_grad=True)

class Connection(nn.Module):
        self.combined_feature_map = torch.zeros([BATCH_SIZE, NUM_OF_IMAGE_CHANNELS, IMAGE_HEIGHT, IMAGE_WIDTH],
                                                requires_grad=True)

def train_NN():
                        # combines all the feature maps from different mixed ops edges
                        graph.cells[c].nodes[n].connections[cc].combined_feature_map += \
                            graph.cells[c].nodes[n].connections[cc].edges[e].forward_f(x)  # Ltrain(w±, alpha)

                    graph.cells[c].nodes[n].output = graph.cells[c].nodes[n].output + \
                        graph.cells[c].nodes[n-1].connections[cc].combined_feature_map

                graph.cells[c].output = graph.cells[c].output + graph.cells[c].nodes[n].output

        outputs1 = graph.cells[NUM_OF_CELLS-1].output