add mb

Author: xadupre

_unittests/ut_helpers/test_torch_test_helper.py

@@ -295,7 +295,14 @@ def forward(self, x, y=None):
         model(x, y)
         stat = model_statistics(model)
         self.assertEqual(
-            {"type": "Model", "n_modules": 1, "param_size": 4, "buffer_size": 4, "float32": 8},
+            {
+                "type": "Model",
+                "n_modules": 1,
+                "param_size": 4,
+                "buffer_size": 4,
+                "float32": 8,
+                "size_mb": 0,
+            },
             stat,
         )
 

onnx_diagnostic/helpers/torch_test_helper.py

@@ -514,8 +514,9 @@ def model_statistics(model: torch.nn.Module):
     res = dict(
         type=model.__class__.__name__,
         n_modules=n_subs,
-        param_size=size,
+        param_size=param_size,
         buffer_size=buffer_size,
+        size_mb=(param_size + buffer_size) // 2**20,
     )
     res.update(sizes)
     return res