Fix reparameterize for NextViT. Fix #2187

rwightman · rwightman · commit 3c0283f9ef93 · 2024-05-27T14:48:58.000-07:00
diff --git a/timm/models/nextvit.py b/timm/models/nextvit.py
@@ -197,7 +197,7 @@ def __init__(
     def reparameterize(self):
         if not self.is_fused:
             merge_pre_bn(self.mlp.fc1, self.norm)
-            self.norm = None
+            self.norm = nn.Identity()
             self.is_fused = True
 
     def forward(self, x):
