From 9e106f34e10bc8ecd3cbdc17d26fa6f1a1bfe728 Mon Sep 17 00:00:00 2001
From: Jared Fernandez <jaredfern@cmu.edu>
Date: Tue, 6 Sep 2022 20:08:27 -0400
Subject: [PATCH] Check for elementwise_affine used by LayerNorm

---
 thop/vision/basic_hooks.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/thop/vision/basic_hooks.py b/thop/vision/basic_hooks.py
index d99a641..fb864b6 100644
--- 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