TypeError: 'Tensor' object is not a mapping

[AKSHILMY]

Describe the issue:

Upon running the following code to compress and speed up the YOLOv5 model using NNI's pruning and speedup techniques.:

import torch

from nni.common.concrete_trace_utils import concrete_trace
from nni.compression.pruning import L1NormPruner
from nni.compression.utils import auto_set_denpendency_group_ids
from nni.compression.speedup import ModelSpeedup

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', device='cpu')

model(torch.rand([1, 3, 640, 640]))

config_list = [{
    'sparsity': 0.5,
    'op_types': ['Conv2d'],
    'exclude_op_names_re': ['model.model.model.24.*'],  # this layer is detector head
}]

config_list = auto_set_denpendency_group_ids(model, config_list, torch.rand([1, 3, 640, 640]))

pruner = L1NormPruner(model, config_list)
masked_model, masks = pruner.compress()
pruner.unwrap_model()

graph_module = concrete_trace(model, (torch.rand([1, 3, 640, 640])))
ModelSpeedup(model, torch.rand([1, 3, 640, 640]), masks, graph_module=graph_module).speedup_model()

model(torch.rand([1, 3, 640, 640]))

The observed error:

/usr/local/lib/python3.10/dist-packages/nni/compression/speedup/utils.py:32: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.
  _register_pytree_node(immutable_dict, _idict_flatten, _idict_unflatten)
/usr/local/lib/python3.10/dist-packages/torch/utils/_pytree.py:254: UserWarning: <class 'torch.fx.immutable_collections.immutable_dict'> is already registered as pytree node. Overwriting the previous registration.
  warnings.warn(
/usr/local/lib/python3.10/dist-packages/nni/compression/speedup/utils.py:33: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.
  _register_pytree_node(immutable_list, _ilist_flatten, _ilist_unflatten)
/usr/local/lib/python3.10/dist-packages/torch/utils/_pytree.py:254: UserWarning: <class 'torch.fx.immutable_collections.immutable_list'> is already registered as pytree node. Overwriting the previous registration.
  warnings.warn(
/usr/local/lib/python3.10/dist-packages/torch/hub.py:294: UserWarning: You are about to download and run code from an untrusted repository. In a future release, this won't be allowed. To add the repository to your trusted list, change the command to {calling_fn}(..., trust_repo=False) and a command prompt will appear asking for an explicit confirmation of trust, or load(..., trust_repo=True), which will assume that the prompt is to be answered with 'yes'. You can also use load(..., trust_repo='check') which will only prompt for confirmation if the repo is not already trusted. This will eventually be the default behaviour
  warnings.warn(
Downloading: "https://github.com/ultralytics/yolov5/zipball/master" to /root/.cache/torch/hub/master.zip
requirements: Ultralytics requirement ['gitpython>=3.1.30'] not found, attempting AutoUpdate...
Collecting gitpython>=3.1.30
  Downloading GitPython-3.1.42-py3-none-any.whl (195 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 195.4/195.4 kB 2.1 MB/s eta 0:00:00
Collecting gitdb<5,>=4.0.1 (from gitpython>=3.1.30)
  Downloading gitdb-4.0.11-py3-none-any.whl (62 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 62.7/62.7 kB 20.4 MB/s eta 0:00:00
Collecting smmap<6,>=3.0.1 (from gitdb<5,>=4.0.1->gitpython>=3.1.30)
  Downloading smmap-5.0.1-py3-none-any.whl (24 kB)
Installing collected packages: smmap, gitdb, gitpython
Successfully installed gitdb-4.0.11 gitpython-3.1.42 smmap-5.0.1

requirements: AutoUpdate success ✅ 6.3s, installed 1 package: ['gitpython>=3.1.30']
requirements: ⚠️ Restart runtime or rerun command for updates to take effect

YOLOv5 🚀 2024-3-25 Python-3.10.12 torch-2.2.1+cu121 CPU

Downloading https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt to yolov5s.pt...
100%|██████████| 14.1M/14.1M [00:00<00:00, 134MB/s]

Fusing layers... 
YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients, 16.4 GFLOPs
Adding AutoShape... 
/root/.cache/torch/hub/ultralytics_yolov5_master/models/yolo.py:100: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if self.dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
/root/.cache/torch/hub/ultralytics_yolov5_master/models/yolo.py:100: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if self.dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
/root/.cache/torch/hub/ultralytics_yolov5_master/models/yolo.py:100: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if self.dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
/root/.cache/torch/hub/ultralytics_yolov5_master/models/yolo.py:100: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if self.dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
[<ipython-input-3-667814b03cd6>](https://localhost:8080/#) in <cell line: 24>()
     22 pruner.unwrap_model()
     23 
---> 24 graph_module = concrete_trace(model, (torch.rand([1, 3, 640, 640])))
     25 ModelSpeedup(model, torch.rand([1, 3, 640, 640]), masks, graph_module=graph_module).speedup_model()
     26 

1 frames
[/usr/local/lib/python3.10/dist-packages/nni/common/concrete_trace_utils/concrete_tracer.py](https://localhost:8080/#) in trace(self, root, autowrap_modules, autowrap_leaf_function, autowrap_leaf_class, leaf_module, fake_middle_class, concrete_args, use_operator_patch, operator_patch_backlist, forward_function_name)
    677         else:
    678             kv_default = {k: v for k, v in zip(args[-len(defaults):], defaults)}
--> 679             concrete_args = {
    680                 **concrete_args,
    681                 **{n: kv_default[n] for n in args if n not in concrete_args}

TypeError: 'Tensor' object is not a mapping

Environment:

• NNI version: 3.0
• Training service (local|remote|pai|aml|etc): remote
• Client OS:
• Server OS (for remote mode only): Google Colab
• Python version: 3.10.12
• PyTorch/TensorFlow version: 2.2.1+cu121
• Is conda/virtualenv/venv used?: No
• Is running in Docker?: No

Configuration:

• Experiment config (remember to remove secrets!):
• Search space:

Log message:

• nnimanager.log:
• dispatcher.log:
• nnictl stdout and stderr:

How to reproduce it?: