I’ve included some sample code below. I am trying to mask the gradients of layer1 in this example. As you can see, I’m zeroing out the weights upon initialisation.
class MaskedLinear(nn.Module):
    # Currently unused. Intended for backward hook.
    def _zero_grad_mask(self, module, grad_input, grad_output):
        new_grad = Variable(
            torch.Tensor(grad_input[0].shape).cuda().masked_scatter_(self.indices_mask_tensor, grad_input[0].data)
        )
        return (grad_input[0], None, new_grad)

    def __init__(self, in_features, out_features, indices_mask):
        super(MaskedLinear, self).__init__()
        self.linear = nn.Linear(in_features, out_features).cuda()
        good_weights = self.linear.weight.data[indices_mask]
        self.linear.weight.data[:] = 0 # zero all weights out
        self.linear.weight.data[indices_mask] = good_weights # populate good weights
        self.indices_mask_tensor = torch.ByteTensor(out_features, in_features).cuda()
        self.indices_mask_tensor[indices_mask] = 1

    def forward(self, input):
        # Repeating original masking here is currently used, and works as a hack, but is incredibly slow.
        good_weights = self.linear.weight.data[self.indices_mask]
        self.linear.weight.data[:] = 0
        self.linear.weight.data[self.indices_mask] = good_weights
        return self.linear(input)

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.layer1 = MaskedLinear(in_dim, out_dim, mask)
        
    def forward(self, x):
        x = F.relu(self.layer1(x))
        return x