Does Torch Pruning support pruning of any model?

[24367452]

I want to check the pruning status through a simple model, but I encountered a dimension mismatch issue. What should I do?

class SimpleModel(nn.Module):     
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.conv1 = nn.Conv1d(in_channels=1, out_channels=128, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(in_channels=128, out_channels=64, kernel_size=3, padding=1)
        self.conv3 = nn.Conv1d(in_channels=64, out_channels=32, kernel_size=3, padding=1)
        self.flatten = nn.Flatten()
        self.dense = nn.Linear(in_features=32*4, out_features=2)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = self.flatten(x)
        x = self.dense(x)
        return x

Report an error:

  File "C:\Users\Administrator\.conda\envs\www\lib\site-packages\torch\nn\functional.py", line 1848, in linear
    return torch._C._nn.linear(input, weight, bias)
RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x112 and 124x1)

What should I do?

[lxaw]

You may be pruning the final classifer, see here: https://github.com/VainF/Torch-Pruning/blob/master/examples/notebook/1%20-%20Customize%20Your%20Own%20Pruners.ipynb