[nn_layer_stats] Fix args passing of layer.total_ifmap/ofmap_size().

stanford-mast · Aug 24, 2020 · 198a527 · 198a527
1 parent d7b7ac3
commit 198a527
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nn_dataflow/tools/nn_layer_stats.py b/nn_dataflow/tools/nn_layer_stats.py
@@ -68,8 +68,8 @@ def layer_stats(args):
         elif isinstance(layer, ConvLayer):
             convs += 1
 
-        ifmap_size = layer.total_ifmap_size(word_bytes) * batch / KILO
-        ofmap_size = layer.total_ofmap_size(word_bytes) * batch / KILO
+        ifmap_size = layer.total_ifmap_size(batch, word_bytes) / KILO
+        ofmap_size = layer.total_ofmap_size(batch, word_bytes) / KILO
         try:
             filter_size = layer.total_filter_size(word_bytes) / KILO
         except AttributeError: