Unable to backpropagate the network with custom layer

ZeeshanH: <p>My network</p>
<pre><code class="lang-auto">class Model(torch.nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        self.mel = Spectrogram.MelSpectrogram(n_fft=512, n_mels=hp.data.nmels, sr=16000, hop_length=160, trainable_mel=True, trainable_STFT=True, device=device) # Using device='cuda:0'
#         torch.nn.init.xavier_uniform_(self.mel.mel_basis)
        self.LSTM_stack = nn.LSTM( hp.data.nmels ,hp.model.hidden , num_layers=hp.model.num_layer, batch_first=True).cuda(device)
        for name, param in self.LSTM_stack.named_parameters():
          if 'bias' in name:
             nn.init.constant_(param, 0.0)
          elif 'weight' in name:
             nn.init.xavier_normal_(param)
        self.CNN_layer1 = torch.nn.Conv2d(161,1, kernel_size=(1,1)).cuda(device)
        self.projection = nn.Linear(hp.model.hidden, hp.model.proj).cuda(device)
        self.clustering=spectral_clusterer.CustomClusteringLayer()

    def forward(self, x,y):
        x = self.mel(x)
        x = x.view(-1,x.data.size()[2],x.data.size()[1])
        x,_=self.LSTM_stack(x)
        x = x.view(-1,x.data.size()[1],x.data.size()[2],1)
        x=self.CNN_layer1(x)
        x=x.squeeze()
        x=self.projection(x)
        x = x / torch.norm(x, dim=1).unsqueeze(1)
        # x=x.squeeze()
        x=self.clustering(x,y)
        return x
</code></pre>
<p>spectral_clusterer.CustomClusteringLayer() has no parameter to learn it performs certain computation. feed forward is okay but usually get error on loss.backward()</p>
<p>This custom layer is defined as.</p>
<pre><code class="lang-auto">class SpectralClusterer(torch.autograd.Function):
      
    @staticmethod
    def forward(ctx,X,clusters):
            gaussian_blur_sigma=1
            p_percentile=0.95
            thresholding_soft_multiplier=0.01
            stop_eigenvalue=1e-2

            if name == "CropDiagonal":
                return refinement.CropDiagonal_torch()
            elif name == "GaussianBlur":
                return refinement.GaussianBlur_torch(gaussian_blur_sigma)
            elif name == "RowWiseThreshold":
                return refinement.RowWiseThreshold_torch(
                    p_percentile,
                    thresholding_soft_multiplier)
            elif name == "Symmetrize":
                return refinement.Symmetrize_torch()
            elif name == "Diffuse":
                return refinement.Diffuse_torch()
            elif name == "RowWiseNormalize":
                return refinement.RowWiseNormalize_torch()
            else:
                raise ValueError("Unknown refinement operation: {}".format(name))
        k=len(clusters)
        if not torch.is_tensor(X):
            raise TypeError("X must be a torch tensor")
        if len(X.shape) != 2:
            raise ValueError("X must be 2-dimensional")
        
        affinity = utils.compute_affinity_matrix_torch(X)
        refinement_sequence=DEFAULT_REFINEMENT_SEQUENCE
        
        for refinement_name in refinement_sequence:
            op = _get_refinement_operator(refinement_name)
            affinity = op.refine(affinity)
        
        cluster_ids_x, cluster_centers = KMeans(X=affinity, num_clusters=k, distance='euclidean', device=torch.device('cuda:0'))
        ctx.save_for_backward(cluster_ids_x,clusters)
        return cluster_ids_x
     
class CustomClusteringLayer(nn.Module):
    def __init__(self,name=""):
        super(CustomClusteringLayer, self).__init__()
        self.name=name
    def forward(self,x,n):
         return SpectralClusterer.apply(x,n)

</code></pre>
<p>Is there any way to bypass backprop in this layer?</p>

ptrblck: <p>You would need to implement the <code>backward</code> method manually for this custom <code>autograd.Function</code>, as Autograd doesn’t know how the gradients should be passed through this layer.<br>
This can be avoided if you could use pure PyTorch methods in the <code>forward</code> pass, but this doesn’t seem to be the case, as you are apparently using methods from other libraries.</p>
<p><a href="https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html">This tutorial</a> gives you an example of a custom function.</p>

ZeeshanH: <p>I followed the tutorial which you mentioned. Moreover all libraries which I am using are Pytorch based. Nothing is implemented which is not in torch. This layer does not have any Parameter to learn then What should I implement in Backward function? I tried returning None in backward but it didnt solve the error.</p>

ptrblck: <p>If all operations in the <code>forward</code> method are differentiable, you could create a custom <code>nn.Module</code> and let Autograd create the backward pass.<br>
<code>autograd.Function</code>s are used for custom forward and backward implementations.</p>

ZeeshanH: <p>Where can I find list of operators which dont break differentiability of torch Graph?</p>

ptrblck: <p>The majority of PyTorch operations is differentiable. I’m not sure if there is a list of non-differentiable operations, but in doubt you could check the <code>.grad_fn</code> of the output.<br>
If it points to a valid function, then Autograd will be able to propely backpropagate through it.</p>
<p>Note that recreating new tensors (e.g. via <code>x = torch.tensor(x)</code>) will break the computation graph.</p>