Improve distance cutoff in DimeNet (#4562)

* Integrate clamping into Envelope

This is cleaner than (dist / cutoff).clamp_max(1.0) because it's the envelope that should become zero and then effectively mask distances larger than the cutoff. This is also how it's implemented in [`DimeNet`](https://github.com/gasteigerjo/dimenet/blob/09123a0e16e728d0a0e53e6686b04f859802aa81/dimenet/model/layers/envelope.py#L23).

* Remove extra parentheses

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update torch_geometric/nn/models/dimenet.py

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Matthias Fey <matthias.fey@tu-dortmund.de>

torch_geometric/nn/models/dimenet.py

@@ -45,7 +45,8 @@ def forward(self, x):
         x_pow_p0 = x.pow(p - 1)
         x_pow_p1 = x_pow_p0 * x
         x_pow_p2 = x_pow_p1 * x
-        return 1. / x + a * x_pow_p0 + b * x_pow_p1 + c * x_pow_p2
+        return (1. / x + a * x_pow_p0 + b * x_pow_p1 +
+                c * x_pow_p2) * (x < 1.0).to(x.dtype)
 
 
 class BesselBasisLayer(torch.nn.Module):
@@ -64,7 +65,7 @@ def reset_parameters(self):
         self.freq.requires_grad_()
 
     def forward(self, dist):
-        dist = (dist.unsqueeze(-1) / self.cutoff).clamp(max=1.0)
+        dist = (dist.unsqueeze(-1) / self.cutoff)
         return self.envelope(dist) * (self.freq * dist).sin()