From acf7d56e363d74e469adb1a09b5fd7b307c5a338 Mon Sep 17 00:00:00 2001
From: Necolizer <58028682+Necolizer@users.noreply.github.com>
Date: Tue, 15 Oct 2024 02:48:05 +0800
Subject: [PATCH] fix bug

fix the AttributeError raised in function _einsum_flops_compute() when using torch.einsum
---
 calflops/pytorch_ops.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/calflops/pytorch_ops.py b/calflops/pytorch_ops.py
index b52b9fa..f95bede 100644
--- a/calflops/pytorch_ops.py
+++ b/calflops/pytorch_ops.py
@@ -292,7 +292,7 @@ def _einsum_flops_compute(equation, *operands):
     Count flops for the einsum operation.
     """
     equation = equation.replace(" ", "")
-    input_shapes = [o.shape for o in operands]
+    input_shapes = [o_element.shape for o in operands for o_element in (o if isinstance(o, list) else [o])]
 
     # Re-map equation so that same equation with different alphabet
     # representations will look the same.