text: import torch
  
  
  @torch.jit.script
  def _poly_cutoff(x: torch.Tensor, factor: float) -> torch.Tensor:
      p: float = 6.0
      x = x * factor
  
      out = 1.0
      out = out - (((p + 1.0) * (p + 2.0) / 2.0) * torch.pow(x, p))
      out = out + (p * (p + 2.0) * torch.pow(x, p + 1.0))
      out = out - ((p * (p + 1.0) / 2) * torch.pow(x, p + 2.0))
  
      return out * (x < 1.0)
  
  
  class PolynomialCutoff(torch.nn.Module):
      _factor: float
      p: float
  
      def __init__(self, r_max: float, p: float = 6):
          r"""Polynomial cutoff, as proposed in DimeNet: https://arxiv.org/abs/2003.03123
  
  
          Parameters
          ----------
          r_max : float
              Cutoff radius
  
          p : int
              Power used in envelope function
          """
          super().__init__()
          if p != 6:
              raise NotImplementedError(
                  "p values other than 6 not currently supported for simplicity; if you need this please file an issue."
              )
          self.p = p
          self._factor = 1.0 / r_max
  
      def forward(self, x):
          """
          Evaluate cutoff function.
  
          x: torch.Tensor, input distance
          """
          return _poly_cutoff(x, self._factor)
  