Fix setattr issue in canonizers.MergeBatchNorm (#208)

chr5tphr · Jul 19, 2024 · 0ace8af · 0ace8af
1 parent e5699aa
commit 0ace8af
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/zennit/canonizers.py b/src/zennit/canonizers.py
@@ -125,8 +125,8 @@ def merge_batch_norm(modules, batch_norm):
 
         for module in modules:
             if module.bias is None:
-                object.__setattr__(
-                    module, 'bias', torch.zeros(1, device=module.weight.device, dtype=module.weight.dtype)
+                setattr(
+                   module, 'bias', torch.nn.Parameter(torch.zeros(1, device=module.weight.device, dtype=module.weight.dtype))
                 )
 
             index = (slice(None), *((None,) * (module.weight.ndim - 1)))