How to generate source code of compressed model with mask?

As far as I know, most channel compression methods generate a mask for each module in terms of channels. For example, we have a model defined by

class Net(nn.Module):
    def __init__(self):
        self.conv1 = nn.Conv2d(3, 10, 3)
        self.conv2 = nn.Conv2d(10, 20, 3)
        self.fc = nn.Linear(20,10)

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.fc(out)
        return out

Suppose we generate a mask as below

{
    'conv1': 5,
    'conv2': 12,
}

The mask means that the output channel size of conv1 and conv2 can be compressed to 5 and 12. Now I know how to generate the mask, but I wonder Is any beautiful and simple method to generate the new code

class Net(nn.Module):
    def __init__(self):
        self.conv1 = nn.Conv2d(3, 5, 3)
        self.conv2 = nn.Conv2d(5, 12, 3)
        self.fc = nn.Linear(12, 10)

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.fc(out)
        return out

Besides, if the source code builds modules by functions (e.g., _make_layer in resnet), can we still generate the code based on the mask?