How to put tensor from linear to conv layer

Hello i have this example, how to reshape output from linear layer if i want to put it inside conv2d ?

import torch
import torch.nn as nn


x = torch.rand((32,1,28,28))
real_inputs = x.view(-1, 784)

linear = nn.Linear(784, 784)
conv = nn.Conv2d(784, 64, kernel_size=9, stride=1, padding=9 // 2)

x = linear(real_inputs)
x = conv(x)
RuntimeError: Expected 4-dimensional input for 4-dimensional weight [64, 784, 9, 9], but got 2-dimensional input of size [32, 784] instead

You would need to add the missing spatial sizes.
The output of linear would have the shape [batch_size=32, features=784].
The conv layer expects an input with 784 channels, so if you don’t want to repeat the activation you would only be able to unsqueeze the missing dimensions since reshaping won’t work either without changing the batch dimension:

x = torch.rand((32,1,28,28))
real_inputs = x.view(-1, 784)

linear = nn.Linear(784, 784)
conv = nn.Conv2d(784, 64, kernel_size=9, stride=1, padding=9 // 2)

x = linear(real_inputs)
x = x[:, :, None, None]

out = conv(x)
print(out.shape)
# > torch.Size([32, 64, 1, 1])
1 Like