fix graph_net/torch/count_kernels_util.py

roll-away · roll-away · commit 7ef1e38fa8e1 · 2025-12-23T07:08:47.000Z
diff --git a/graph_net/torch/count_kernels_util.py b/graph_net/torch/count_kernels_util.py
@@ -39,7 +39,7 @@ def run():
 
     # Use PyTorch Profiler
     with profile(
-        activities=[ProfilerActivity.CUDA],
+        activities=[ProfilerActivity.CUDA, ProfilerActivity.CPU],
         record_shapes=True,
     ) as prof:
         with record_function("model_inference"):
