Torch-Pruning KeyError while attempting to prune yolov5 model

Hello,

I am trying to prune yolov5 model using Torch_Pruning. It fails with error message: KeyError: Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))

Detailed Traceback: Traceback (most recent call last): File "models/prune_TP_git_issue.py", line 117, in new_model = prune_model(model, img) File "models/prune_TP_git_issue.py", line 64, in prune_model prune_conv(mm3, SPARSITY) File "models/prune_TP_git_issue.py", line 48, in prune_conv raise e File "models/prune_TP_git_issue.py", line 46, in prune_conv plan = DG.get_pruning_plan(conv, tp.prune_conv, prune_index) File "/home/tfs/venv_yolov5_Torch_Pruning/lib/python3.6/site-packages/torch_pruning-0.2.2-py3.6.egg/torch_pruning/dependency.py", line 330, in get_pruning_plan KeyError: Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))

Attaching the prune_TP.py code. (a) save the code under models dir (b) download pre-trained model and place it under weights dir (c) invoke the code from the main dir as follows: $ python models/prune_TP_git_issue.py

Also attaching the logfile from this run. It looks like DG.module_to_node dictionary is not built correctly - it seems to have only 1 entry: ------------ BEGIN : DG.module_to_node --------------------- {Conv2d(512, 255, kernel_size=(1, 1), stride=(1, 1)): <Node: (model.24.m.2 (Conv2d(512, 255, kernel_size=(1, 1), stride=(1, 1))), None)>} ------------ END : DG.module_to_node ---------------------

prune_TP_git_issue.yolov5.py.zip log.prune_TP_git_issue.txt