Here, also, is my implementation:

  
    
    
    Gradient masking in register_backward_hook for custom connectivity - efficient implementation 
  
  
    @jonathanEphrath Here is my implementation which was inspired by @jukiewiczm’s code. 
class MaskedLinear(nn.Module):
    def __init__(self, in_dim, out_dim, indices_mask):
        """
       :param in_features: number of input features
       :param out_features: number of output features
       :param indices_mask: list of two lists containing indices for dimensions 0 and 1, used to create the mask
       """
        super(MaskedLinear, self).__init__()
 
        def backward_hook(grad):
      …
  


@uchida-takumi I would be interested in your thoughts on comparisons between the approaches.