[Bug] benchmark shows 0.00MB consumed

[deciding]

🐛 Bug

benchmark shows that the memory consumption of some linear/activation layers is 0.00MB

Command

pytest benchmark/benchmark_sp24.py

To Reproduce

Steps to reproduce the behavior:

1.pytest benchmark/benchmark_sp24.py

Expected behavior

It should show the memory consumed during the execution.
But it always shows 0.00MB consumed.

Environment

• PyTorch Version (e.g., 1.0): 2.2.2
• OS (e.g., Linux): Ubuntu 22.04.3 LTS (x86_64)
• How you installed PyTorch (conda, pip, source): pip
• Build command you used (if compiling from source): pip install -e ., pip install -r requirements-benchmark.py
• Python version: 3.11.5
• CUDA/cuDNN version: 12.3.107, /usr/local/cuda-12.3/targets/x86_64-linux/lib/libcudnn.so.8.9.7
• GPU models and configuration: NVIDIA A100-SXM4-80GB

[deciding]

submmited PR #1018

[deciding]

Anyone help look at this issue?