Hi; Consider I have a neural network ;
nn.Sequential(
            nn.Linear(in_features, out_features)
            UnFlatten(), # reshape the output of previous layers so that its input shape to the next layers is proper
            nn.ConvTranspose2d()
            nn.ReLU(),
            nn.ConvTranspose2d()
            nn.ReLU(),
            nn.ConvTranspose2d()
            nn.ReLU(),
            nn.ConvTranspose2d()
            nn.Sigmoid(),
        )


in the regular cases; it takes data X.shape = (batch_size, in_features); and output a batch of images,
but right now; I have an input looks like this X_shape = ( batch_size, another_features, in_features) and I want to output a batch of batch of images
In the for loop way, the implementation goes like this
result = []
for i in range(another_features):
    temp.append(NN(X[:,i,:])) # feed each slice of another_features 
result = torch.stack(result) # leads to 5D tensor has shape (another_features, batch_sizes, out_channels, out_features, out_features)
# then reshape into  (batch_sizes , another_features, out_channels, out_features, out_features)

I’m wondering if there is an efficient way because technically, the process isn’t sequential and can be parallelized.