def forward(self, distances):
        """

        Description
        -----------
        Apply the layer to transform edge distances.

        Parameters
        ----------
        distances : Float32 tensor of shape (E, 1)
            Distance between end nodes of edges. E for the number of edges.

        Returns
        -------
        Float32 tensor of shape (K, E, 1)
            Transformed edge distances. K for the number of radial filters.
        """
        scaled_euclidean_distance = (
            -self.rbf_kernel_scaling * (distances - self.rbf_kernel_means) ** 2
        )  # (K, E, 1)
        rbf_kernel_results = th.exp(scaled_euclidean_distance)  # (K, E, 1)

        cos_values = 0.5 * (
            th.cos(np.pi * distances / self.interaction_cutoffs) + 1
        )  # (K, E, 1)
        cutoff_values = th.where(
            distances <= self.interaction_cutoffs,
            cos_values,
            th.zeros_like(cos_values),
        )  # (K, E, 1)

        # Note that there appears to be an inconsistency between the paper and
        # DeepChem's implementation. In the paper, the scaled_euclidean_distance first
        # gets multiplied by cutoff_values, followed by exponentiation. Here we follow
        # the practice of DeepChem.
        return rbf_kernel_results * cutoff_values


        Description
        -----------
        Apply the layer to transform edge distances.

        Parameters
        ----------
        distances : Float32 tensor of shape (E, 1)
            Distance between end nodes of edges. E for the number of edges.

        Returns
        -------
        Float32 tensor of shape (K, E, 1)
            Transformed edge distances. K for the number of radial filters.
        
python/dgl/nn/pytorch/conv/atomicconv.py