Update torch_geometric/nn/models/dimenet.py

pyg-team · Apr 29, 2022 · 133e195 · 133e195
1 parent 101f754
commit 133e195
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_geometric/nn/models/dimenet.py b/torch_geometric/nn/models/dimenet.py
@@ -46,7 +46,7 @@ def forward(self, x):
         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) * (x < 1.0).float()
+                c * x_pow_p2) * (x < 1.0).to(x.dtype)
 
 
 class BesselBasisLayer(torch.nn.Module):