I’ve set up a function which could possibly help you out.
def revert_mask(maskdict, net, netdict_orig):
    '''
    Reverts certain elements to it's original state.
    Values with TRUE in mask are reverted to original state.
    Values with FALSE in mask are forced to be at 0.
    '''
    for name, param in net.named_parameters():
        mask = maskdict[name]
        param.data[mask] = netdict_orig[name][mask]
        param.data[~mask] = 0.0 # Force values at stay at 0
    return net

def update_grad(net, maskdict):
    for name, param in net.named_parameters():
        mask = maskdict[name]
        param.grad[~mask] = 0 # Zero out grad values.
    return net

Given a mask, it’ll zero out weights (and gradients, if you need to - I’ve kept them in different functions). Hope this helps.