A question about applying a neural network on a tensor's specific dimension

1

Hello I have a little question that might be straightforward:

Suppose we have a torch.tensor x with shape (4,5,1).

What I’d like to do is applying a neural network on this input data, but on the last dimension only. Let’s say we decide some output features having dimension 64, then my output’s shape would be something like (4,5,64).

Trying with the standard procedure, I have noticed that the input tensor gets automatically flattened and the input seeks to receive an input of dimensions (20,1)

Many thanks,

Federico

2

Hi

This can be done but it’ll depend on the network’s architecture, so for example for a simple network with 1 nn.Linear layer. Both the batch dimensions are outputted. For example,

>>> x = torch.randn(4,5,64)
>>> fc = torch.nn.Linear(64, 1)
>>> y = fc(x)
>>> y.shape #torch.Size([4, 5, 1])

Can you share some more about the model?

1 Like