[fx] make const fold code more pythonic (#64451)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/64451
No functional change.
Test Plan:
```
buck test caffe2/test:fx_const_fold
```
Reviewed By: jfix71, RoshanPAN, houseroad
Differential Revision: D30718255
fbshipit-source-id: 95f98561c7f33fcc6c839db68683c85eb152c949