profile_statistic (#233)

Co-authored-by: Cano Xiao <cx922@ee-kraken.ee.ic.ac.uk>
DeepWok · Nov 12, 2024 · 9bf2bcf · 9bf2bcf
1 parent 149fc06
commit 9bf2bcf
Showing 1 changed file with 5 additions and 8 deletions.
diff --git a/src/chop/passes/graph/analysis/statistical_profiler/profile_statistics.py b/src/chop/passes/graph/analysis/statistical_profiler/profile_statistics.py
@@ -201,18 +201,15 @@ def run_node(self, n: Node) -> Any:
             if len(act_stats) > 0:
                 numeric_args = tuple(
                     filter(
-                        lambda x: isinstance(x, (torch.Tensor, int, float))
+                        lambda x: isinstance(x, (torch.Tensor))
                         and not isinstance(x, bool),
                         args + tuple(kwargs.values()),
                     )
                 )
-                try:
-                    assert len(numeric_args) == len(act_stats), (
-                        f"Number of tensor args ({len(numeric_args)}) "
-                        f"does not match number of act entries ({len(act_stats)})"
-                    )
-                except AssertionError:
-                    exit(1)
+                assert len(numeric_args) == len(act_stats), (
+                    f"Number of tensor args ({len(numeric_args)}) "
+                    f"does not match number of act entries ({len(act_stats)})"
+                )
 
                 for tensor_arg, stat in zip(numeric_args, act_stats):
                     stat.update(tensor_arg)