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]

conda install pytorch==2.1.0 torchvision torchaudio cudatoolkit=11.8 -c pytorch -c nvidia

Downgrade the Pytorch version to 2.1.0; it works.