From 18988e19a9f49d3f976f46aacfc046949be37076 Mon Sep 17 00:00:00 2001
From: frgfm <fgfm03@hotmail.fr>
Date: Wed, 3 Mar 2021 17:18:30 +0100
Subject: [PATCH] style: Fixed typing

---
 torchscan/modules/memory.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/torchscan/modules/memory.py b/torchscan/modules/memory.py
index ae8883c..b3082b5 100644
--- a/torchscan/modules/memory.py
+++ b/torchscan/modules/memory.py
@@ -203,7 +203,7 @@ def dmas_pool(module: Union[_MaxPoolNd, _AvgPoolNd], input: Tensor, output: Tens
     # Resolve kernel size and stride size (can be stored as a single integer or a tuple)
     if isinstance(module.kernel_size, tuple):
         kernel_size = module.kernel_size
-    else:
+    elif isinstance(module.kernel_size, int):
         kernel_size = (module.kernel_size,) * (input.ndim - 2)  # type: ignore[attr-defined]
 
     # Each output element required K ** 2 memory accesses