QuickStart example for pruning is not working properly !!

[hamdy-cryptic]

Describe the bug:
QuickStart Example provides Error at model SpeedUp. AttributeError: 'NoneType' object has no attribute 'startswith'. I tried running the pruning example quick start code and all worked fine until the section of model speedup.

Environment:

• NNI version: 3
• Training service (local|remote|pai|aml|etc):
• Python version: 3.10
• PyTorch version: 2.2
• Cpu or cuda version: CPU

Reproduce the problem

• Code|Example:
  from nni.compression.speedup import ModelSpeedup

ModelSpeedup(model, torch.rand(3, 1, 28, 28).to(device), masks).speedup_model()

• How to reproduce:
  AttributeError Traceback (most recent call last)
  in <cell line: 7>()
  5 from nni.compression.speedup import ModelSpeedup
  6
  ----> 7 ModelSpeedup(model, torch.rand(3, 1, 28, 28).to(device), masks).speedup_model()

2 frames
/usr/local/lib/python3.10/dist-packages/nni/common/concrete_trace_utils/concrete_tracer.py 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)
956 elif func.name != func.qualname and func.qualname != 'boolean_dispatch..fn':
957 # method
--> 958 if func.module.startswith('_') and func.module != 'main':
959 path = sys.modules[func.module[1:]]
960 else:

AttributeError: 'NoneType' object has no attribute 'startswith'

[BitCircuit]

Same problem here

[EzcodingSen]

same

[LongBaoCoder2]

using pytorch version 2.1. it works for me

[yating-zh]

!pip install torch torchvision==0.16.0 torchaudio==2.1.0 pytorch-cuda==12.1
Downgrade the Python version to 2.1.0; it works.