Fix bessel GPU (#27)

mir-group · May 4, 2021 · 6de1cb9 · 6de1cb9
1 parent 349b269
commit 6de1cb9
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/nequip/nn/radial_basis.py b/nequip/nn/radial_basis.py
@@ -1,5 +1,6 @@
+import math
+
 import torch
-import numpy as np
 
 from torch import nn
 
@@ -31,12 +32,13 @@ def __init__(self, r_max, num_basis=8, trainable=True):
         self.r_max = float(r_max)
         self.prefactor = 2.0 / self.r_max
 
-        self.bessel_weights = torch.linspace(
-            start=1.0, end=num_basis, steps=num_basis
-        ) * torch.Tensor([np.pi])
-
+        bessel_weights = (
+            torch.linspace(start=1.0, end=num_basis, steps=num_basis) * math.pi
+        )
         if self.trainable:
-            self.bessel_weights = nn.Parameter(self.bessel_weights)
+            self.bessel_weights = nn.Parameter(bessel_weights)
+        else:
+            self.register_buffer("bessel_weights", bessel_weights)
 
     def forward(self, x):
         """