Pass weight as a list:
optimizer = optim.SGD([weight], lr=0.001)


Add all parameters to a list or use the module approach:



class MyConvZ(nn.Module):
    def __init__(self):
        super(MyConvZ, self).__init__()
        self.fn = convZ.apply
        self.weight = nn.Parameter(torch.randn(1, 1, 2, 2))
        
    def forward(self, x):
        x = self.fn(x, self.weight)
        return x

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.convz = MyConvZ()
        self.pool = nn.AvgPool1d(2)

    def forward(self, x):
        x = self.pool(F.relu(self.convz(x)))
        return x

optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

However, the shapes do not match for an image tensor as the input.
The input is expected to have the shape [batch_size, channels, height, width], while the conv weights should be [nb_kernel, in_channels, height, width].