Error when using nn.Upsample in ModelSpeedup #4416
Description
When my model contains nn.Upsample, it reports error. As long as i comment the upsample module, the code runs OK.
It seems that length of dummy_input is 3
My test network:
class TestNet(nn.Module):
def init(self):
super(TestNet, self).init()
self.conv = nn.Conv2d(3, 3, 3, padding=1, stride=1, bias=True)
self.down_sample_1 = nn.Upsample(scale_factor=2)
def forward(self, x):
x = self.down_sample_1(x) # commit this line, the code is OK
x = self.conv(x)
return x
Error Msg:
Traceback (most recent call last):
File "test_pruner_1.py", line 59, in
m_speedup.speedup_model()
File "/home/lvweijie/miniconda3/envs/pt1.5/lib/python3.6/site-packages/nni/compression/pytorch/speedup/compressor.py", line 512, in speedup_model
self.infer_modules_masks()
File "/home/lvweijie/miniconda3/envs/pt1.5/lib/python3.6/site-packages/nni/compression/pytorch/speedup/compressor.py", line 358, in infer_modules_masks
self.update_direct_sparsity(curnode)
File "/home/lvweijie/miniconda3/envs/pt1.5/lib/python3.6/site-packages/nni/compression/pytorch/speedup/compressor.py", line 228, in update_direct_sparsity
state_dict=copy.deepcopy(module.state_dict()), batch_dim=self.batch_dim)
File "/home/lvweijie/miniconda3/envs/pt1.5/lib/python3.6/site-packages/nni/compression/pytorch/speedup/infer_mask.py", line 87, in init
self.output = self.module(*dummy_input)
File "/home/lvweijie/miniconda3/envs/pt1.5/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in call
result = self.forward(*input, **kwargs)
TypeError: forward() takes 2 positional arguments but 4 were given