zero for second residual grad

[aijianiula0601]

Thanks for you jobs. When we checked the code, we found that there was no gradient for residual layer after second layer, please confirm it.

we change the code to : residual = residual - quantized ---> residual = residual - quantized.detach()

Here's the verification we did

    if __name__ == "__main__":
          quantizer = ResidualVQ(
              num_quantizers=4, dim=256, codebook_size=16,
              kmeans_init=True, kmeans_iters=10, threshold_ema_dead_code=2, channel_last=False,
          )
  
          for i in range(4):
              input = torch.rand((2, 256, 30), requires_grad=True)
              quantized, indices, losses = quantizer(input)
              print(quantized.shape, indices.shape, losses.shape)
  
              losses[0, i].backward()
              print(input.grad)

[lucidrains]

@aijianiula0601 i do believe you are correct! thank you for spotting this! 🙏