Make sure conv weight is registered as a param

jpeg_layers/convolution.py

@@ -11,6 +11,8 @@ class Conv2d(torch.nn.modules.Module):
         self.weight = conv_spatial.weight
         self.padding = conv_spatial.padding
 
+        self.register_parameter('weight', self.weight)
+
         self.register_buffer('J', J[0])
         self.register_buffer('J_i', J[1])