If you don’t have skip-layer connections, the order should usually be the same. Just note that you need to prune weight and bias in Conv2d and BatchNorm2d (if you use them).
If you have skip-layer connections, you can build a dependency graph about who passes output to who. You can then use the dependency graph to trace which convs take input from which convs. One example that builds such a graph can be found here 28 (look up pytorch_visualize.py). This example performs graph visualization.
A much simpler way, however, is to “hijack” the forward function of Conv2d, BatchNorm2d, activations functions and so on, like so:
old_forwards = {nn.Conv2d: nn.Conv2d.forward}

def hijack_forward(self, input):
    if hasattr(input, 'drop_mask'):
        drop_input_channels(self, var.drop_mask)  # modify weights and biases
    drop_mask = drop_output_channels(self)        # modify weights and biases
    output = old_forwards[self.__class__](self, input)
    output.drop_mask = drop_mask
    return output

nn.Conv2d.forward = hijack_forward
# perform one forward pass
# reset hijacked forward functions