Unexpected behavior with object.__setattr__ in canonizers.MergeBatchNorm

[karray]

Description

The c30e3cc4586b43a0cff37e8e74e63cfb9324c710 commit breaks saliency map generation. The behavior of object.__setattr__ for setting the bias parameter in canonizers.MergeBatchNorm differs from using direct assignment or setattr() (object.__setattr__ bypasses custom nn.Module.__setattr__). Since the bias is None, this will lead to inconsistencies in parameter registration.

Input	Result	Expected result

Reproducible Example

import torch
import torch.nn as nn

class CustomModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(1))

    def __setattr__(self, name, value):
        super().__setattr__(name, value)

module = CustomModule()

# Method 1: Direct assignment
module.bias1 = nn.Parameter(torch.zeros(1))

# Method 2: Using setattr
setattr(module, 'bias2', nn.Parameter(torch.zeros(1)))

# Method 3: Using object.__setattr__ bias3 won't be registered as a parameter
object.__setattr__(module, 'bias3', nn.Parameter(torch.zeros(1)))

print("Registered parameters:", module._parameters.keys()) # => ['weight', 'bias1', 'bias2']

[chr5tphr]

Hey @karray,

thanks a lot for spotting this, and for the PR. I guess the issue is that we modify the bias, but it remains unused by the module unless it is explicitly registered as a Parameter. As we are using a Canonizer, I think just adding the Parameter is the right call. I will leave you a small comment on the PR.