Point the to right _BatchNorm.

Author: datumbox

references/classification/utils.py

@@ -197,7 +197,7 @@ def bn_reinitialization(model: torch.nn.Module, gamma: float = 1.0, beta: float
         beta (float): The beta initial value.
     """
     for module in model.modules():
-        if isinstance(module, torch.nn._BatchNorm):
+        if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
             module.weight.fill_(gamma)
             module.bias.fill_(beta)