Make sure conv weight is registered as a param

b36dfadd · Max Ehrlich · 81d5c18c · b36dfadd
Verified Commit b36dfadd authored Dec 05, 2018 by Max Ehrlich
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

jpeg_layers/convolution.py jpeg_layers/convolution.py +2 -0

No files found.
--- a/jpeg_layers/convolution.py
+++ b/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])