nni: Pruning model speed up error

After using the pruning API on the model, when I try to perform the speed up, it throws this error :

Error: TypeError: forward() missing 1 required positional argument: input

Traceback: m_speedup.speedup_model()

File “/pruning/lib/python3.6/site-packages/nni/compression/pytorch/speedup/compressor.py”, line 503, in speedup model self.infer_modules_masks() File “/pruning/Python-test/lib/python3.6/site-packages/nni/compression/pytorch/speedup/compressor.py”, line 349. in infer_modules masks self.update_direct_sparsity(curnode) File “/pruning/Python-test/lib/python3.6/site-packages/nni/compression/pytorch/speedup/compressor.py”, line 219, in update direct sparsity state_dict-copy.deepcopy(module.state_dict()). batch_dim=self.batch_dim) File “/pruning/Python-test/lib/python3.6/site-packages/nni/compression/pytorch/speedup/infer_mask.py”, line 80, in_init__ self.output = self.module("dummy_input) File “/pruning/Python-test/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 889, in _call_impl result = self.forward("input, **kwargs)