From e79592a4cb1a8915866ad4deb463c49c66c22795 Mon Sep 17 00:00:00 2001 From: Yosua Michael Maranatha Date: Tue, 24 May 2022 20:09:02 +0100 Subject: [PATCH] Fix bug by checking if norm_layer weight is None before init --- torchvision/models/resnet.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchvision/models/resnet.py b/torchvision/models/resnet.py index 93f0d06f67b..70602705521 100644 --- a/torchvision/models/resnet.py +++ b/torchvision/models/resnet.py @@ -217,9 +217,9 @@ def __init__( # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 if zero_init_residual: for m in self.modules(): - if isinstance(m, Bottleneck): + if isinstance(m, Bottleneck) and m.bn3.weight is not None: nn.init.constant_(m.bn3.weight, 0) # type: ignore[arg-type] - elif isinstance(m, BasicBlock): + elif isinstance(m, BasicBlock) and m.bn2.weight is not None: nn.init.constant_(m.bn2.weight, 0) # type: ignore[arg-type] def _make_layer(