Fix model_flops feature on PT2

[xuzhao9]

Now the flops counter is broken, need to fix it and add a test to prevent future issues.

Run command:

buck2 run @mode/opt //pytorch/benchmark:run -- resnet50 -d cuda -t train --metrics model_flops --torchdynamo inductor

or

python run.py resnet50 -d cuda -t train --metrics model_flops  --torchdynamo inductor

[xuzhao9]

We should also add a test internally to guard on this use case.