Hello
I would like to do learning weights of a convolution but using torch.nn.functional.conv2d instead of torch.nn.conv2d as in this tutorial: https://pytorch.org/tutorials/beginner/blitz/neural_networks_tutorial.html 2 , because the weights tensor when using torch.nn.conv2d has the shape (C_out, C_in, H_in, H_w) and I rather want (C_out, C_in, 1,1) in my model.
%matplotlib inline
import torch
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
#         self.w_i = torch.ones(10, 3, 1, 1, requires_grad=True)
        self.w_i = nn.Parameter(torch.FloatTensor(3,3,1,1))
   
    def forward(self, x):
        x = F.conv2d(x, self.w_i)
        return x

net = Net()
print(net)

input = torch.ones(1, 3, 5, 5)
true_weights = torch.Tensor([[[[ 2.]],[[ 3.]],[[ 1.]]],
[[[ 0.]],[[ 5.]],[[ 2.]]],
[[[ 1.]],[[ 1.]],[[ 2.]]]])
target = F.conv2d(input,true_weights)
criterion = nn.MSELoss()

import torch.optim as optim
optimizer = optim.SGD(net.parameters(), lr=0.01)

for epoch in range(10000):
    optimizer.zero_grad()   # zero the gradient buffers
    output = net(input)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()    # Does the update

list(net.parameters())

I am new on Pytorch and the code above returns me weights in net.parameters() that is not correct (I want true_weights).
What should I modify to make the code run ?