Error in pruner.step() aufter update to 1.1.9

Hi,

I updated from 1.0.0 to 1.1.9 today and now I get the following error when running pruner.step():

 pruner.step()
  File "/home/mle/prune_data_collection/pruning_benchmark_tool/.venv/lib/python3.8/site-packages/torch_pruning/pruner/algorithms/metapruner.py", line 159, in step
    for group in self.prune_global():
  File "/home/mle/prune_data_collection/pruning_benchmark_tool/.venv/lib/python3.8/site-packages/torch_pruning/pruner/algorithms/metapruner.py", line 257, in prune_global
    imp = torch.cat([local_imp[-1]
RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 512 but got size 256 for tensor number 1 in the list.

I downgraded the version to check back and my code still works on 1.0.0. This is the setup I use for pruning a ResNet trained on CIFAR10:

criterion = tp.importance.MagnitudeImportance(p=2, group_reduction=None)

ignored_layers = []
for m in model.modules():
    if isinstance(m, torch.nn.Linear) and m.out_features == 10:
        ignored_layers.append(m)

# Pruner initialization
iterative_steps = 20
pruner = tp.pruner.MagnitudePruner(
    model,
    example_inputs,
    global_pruning=True,
    importance=criterion,
    iterative_steps=iterative_steps,
    ch_sparsity=0.9,
    ignored_layers=ignored_layers,
)

Currently, I'm fine with running 1.0.0 so I've not spent time on debugging it but I might do so later. Maybe it's just a minor thing and easy to change.

Thanks

VainF / Torch-Pruning

Error in pruner.step() aufter update to 1.1.9 #208