hi, i currently have a model that intakes a tensor N x C x H x W and a label tensor N x L. In the forward method, I want to first use different net for different label and then combine them together for other layers. so i current have some code like:
def forward(self, data):
    for i in range(label_types_num):                                                  
        idx = get_label_idx(i) # get all the index that has this label
        group_data = data[idx, :]                                               
        o = self.seperate_nets[i](group_data)                                                                                                
        output_tensor[idx, :] = o                                               
    # then output_tensor is passed to other layers...             


Just wondering is there a way to parallel do the for loop?