Make sure conv weight is registered as a param

parent 81d5c18c
......@@ -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])
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment