CODEWITHSUNDEEP
pytorch
James Arten
James Arten

Reputation: 666

A question about applying a neural network on a specified dimension using PyTorch

I'm wondering about how to do the following thing:

If I have a torch.tensor x with shape (4,5,1) how can apply a neural network using PyTorch on the last dimension?

Using the standard procedure, the model is flattening the entire tensor into some new tensor of shape (20,1) but this is not actually what I want.

Let's say we want some output features of dimension 64, then I would like to obtain a new object of shape (4,5,64)

Upvotes: 0

Views: 28

Answers (1)

ricagj
ricagj

Reputation: 11

import torch
import torch.nn as nn

x = torch.randn(4, 5, 1)
print(x.size())

# https://pytorch.org/docs/stable/generated/torch.nn.Linear.html
m = nn.Linear(1, 64)

y = m(x)
print(y.size())

result:

torch.Size([4, 5, 1])
torch.Size([4, 5, 64])

Upvotes: 1

Related Questions