Merge pull request #72 from Mrliduanyang/patch-1

Update torchstat_utils.py

Author: sytelus

tensorwatch/model_graph/torchstat_utils.py

@@ -73,7 +73,7 @@ def model_stats2df(model_stats:ModelStats):
 
     df.rename(columns={'name': 'module name',
                        'input_shape': 'input shape',
-                       'input_shape': 'input shape',
+                       'output_shape': 'output shape',
                        'inference_memory': 'infer memory(MB)',
                        'mread': 'MemRead(B)',
                        'mwrite': 'MemWrite(B)'
@@ -87,4 +87,4 @@ def model_stats2df(model_stats:ModelStats):
     #summary += "Total MAdd: {}MAdd\n".format(_round_value(total_operation_quantity))
     #summary += "Total Flops: {}Flops\n".format(_round_value(total_flops))
     #summary += "Total MemR+W: {}B\n".format(_round_value(total_memrw, True))
-    return df
+    return df