[FIXED] Partially stacking Tensors

Issue

I have some output from a convolutional layer in PyTorch of the shape [b,c,h] where b is my batches, c is my channels, and h is the features. I want to stack these to feed into a fully-connected layer without changing the batches, so that they are in the shape [b, c*h]. How can I do this?

Solution

Seems like a simple reshape or view should work:

input.shape # [b,c,h]

reshaped = input.view(input.shape[0],-1) # [b,c*h]

Answered By - DerekG