Merge pull request Lyken17#189 from JaredFern/master

Check for elementwise_affine used by LayerNorm
ultralytics · Sep 7, 2022 · 43c064a · 43c064a
2 parents 54e60ed + 9e106f3
commit 43c064a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/thop/vision/basic_hooks.py b/thop/vision/basic_hooks.py
@@ -64,7 +64,7 @@ def count_normalization(m: nn.modules.batchnorm._BatchNorm, x, y):
     x = x[0]
     # bn is by default fused in inference
     flops = calculate_norm(x.numel())
-    if m.affine:
+    if (getattr(m, 'affine', False) or getattr(m, 'elementwise_affine', False)):
         flops *= 2
     m.total_ops += flops