profile_macs error at PyTorch 1.4

[kim-sunghoon]

Hi, Thank you for the awesome tool! profile_macs is not working at PyTorch 1.4 The key error message is

AttributeError: module 'torch.jit' has no attribute 'get_trace_graph'

Thank you in advance!

Here is my entire error log:

Traceback (most recent call last):
  File "gen_summary.py", line 54, in <module>
    macs = profile_macs(model, inputs)
  File "/home/shkim/.conda/envs/torch1.4/lib/python3.6/site-packages/torchprofile/profile.py", line 12, in profile_macs
    graph = trace(model, args, kwargs)
  File "/home/shkim/.conda/envs/torch1.4/lib/python3.6/site-packages/torchprofile/utils/trace.py", line 17, in trace
    trace, _ = torch.jit.get_trace_graph(Flatten(model), args, kwargs)
AttributeError: module 'torch.jit' has no attribute 'get_trace_graph'

my env setting as follows:

# packages in environment at **:
#
# Name                    Version                   Build  Channel
_libgcc_mutex             0.1                        main  
_tflow_select             2.1.0                       gpu  
absl-py                   0.8.1                    py36_0  
astor                     0.8.0                    py36_0  
atomicwrites              1.3.0                    py36_1  
attrs                     19.3.0                     py_0  
blas                      1.0                         mkl  
brevitas                  0.2.0a0                  pypi_0    pypi
c-ares                    1.15.0            h7b6447c_1001  
ca-certificates           2019.11.27                    0  
certifi                   2019.11.28               py36_0  
cudatoolkit               10.1.243             h6bb024c_0  
cudnn                     7.6.5                cuda10.1_0  
cupti                     10.1.168                      0  
docrep                    0.2.7                    pypi_0    pypi
freetype                  2.8                  hab7d2ae_1  
gast                      0.3.2                      py_0  
google-pasta              0.1.8                      py_0  
grpcio                    1.14.1           py36h9ba97e2_0  
h5py                      2.9.0            py36h7918eee_0  
hdf5                      1.10.4               hb1b8bf9_0  
importlib_metadata        1.3.0                    py36_0  
intel-openmp              2019.4                      243  
jpeg                      9b                   h024ee3a_2  
keras-applications        1.0.8                      py_0  
keras-preprocessing       1.1.0                      py_1  
libgcc-ng                 9.1.0                hdf63c60_0  
libgfortran-ng            7.3.0                hdf63c60_0  
libpng                    1.6.37               hbc83047_0  
libprotobuf               3.11.2               hd408876_0  
libstdcxx-ng              9.1.0                hdf63c60_0  
libtiff                   4.1.0                h2733197_0  
markdown                  3.1.1                    py36_0  
mkl                       2019.4                      243  
mkl-service               2.3.0            py36he904b0f_0  
mkl_fft                   1.0.15           py36ha843d7b_0  
mkl_random                1.1.0            py36hd6b4f25_0  
more-itertools            8.0.2                      py_0  
natsort                   7.0.0                    pypi_0    pypi
ninja                     1.9.0            py36hfd86e86_0  
numpy                     1.18.1           py36h4f9e942_0  
numpy-base                1.18.1           py36hde5b4d6_0  
olefile                   0.46                       py_0  
openssl                   1.0.2u               h7b6447c_0  
packaging                 20.0                       py_0  
pillow                    4.2.1            py36h9119f52_0  
pip                       19.3.1                   py36_0  
pluggy                    0.13.1                   py36_0  
protobuf                  3.11.2           py36he6710b0_0  
py                        1.8.1                      py_0  
pyparsing                 2.4.6                      py_0  
pytest                    5.0.1                    py36_0  
python                    3.6.0                         0  
pytorch                   1.4.0           py3.6_cuda10.1.243_cudnn7.6.3_0    pytorch
readline                  6.2                           2  
scipy                     1.3.2            py36h7c811a0_0  
setuptools                44.0.0                   py36_0  
six                       1.13.0                   py36_0  
sqlite                    3.13.0                        0  
tensorboard               1.14.0           py36hf484d3e_0  
tensorboardx              2.0                      pypi_0    pypi
tensorflow                1.14.0          gpu_py36h3fb9ad6_0  
tensorflow-base           1.14.0          gpu_py36he45bfe2_0  
tensorflow-estimator      1.14.0                     py_0  
tensorflow-gpu            1.14.0               h0d30ee6_0  
termcolor                 1.1.0                    py36_1  
tk                        8.5.18                        0  
torchprofile              0.0.1                    pypi_0    pypi
torchsummary              1.5.1                    pypi_0    pypi
torchvision               0.5.0                py36_cu101    pytorch
tqdm                      4.41.1                   pypi_0    pypi
wcwidth                   0.1.7                    py36_0  
werkzeug                  0.16.0                     py_0  
wheel                     0.33.6                   py36_0  
wrapt                     1.11.2           py36h7b6447c_0  
xz                        5.2.4                h14c3975_4  
zipp                      0.6.0                      py_0  
zlib                      1.2.11               h7b6447c_3  
zstd                      1.3.7                h0b5b093_0  

[zhijian-liu]

Hi, I've just updated the tool to support PyTorch 1.4.0. Please let me know if you have any further questions.

[xieydd]

Works for 1.4.0