PyTorch Geometric custom layer parameters not updating

dense_to_sparse does not break the gradient.
corr_layer not updated at all?