add cpu() to avoid crash when using ase with GPU (#156)

materialsvirtuallab · Sep 1, 2023 · e825b75 · e825b75
1 parent f1e2efe
commit e825b75
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/matgl/ext/ase.py b/matgl/ext/ase.py
@@ -177,14 +177,14 @@ def calculate(
         else:
             energies, forces, stresses, hessians = self.potential(graph, state_attr_default)
         self.results.update(
-            energy=energies.detach().numpy(),
-            free_energy=energies.detach().numpy(),
-            forces=forces.detach().numpy(),
+            energy=energies.detach().cpu().numpy(),
+            free_energy=energies.detach().cpu().numpy(),
+            forces=forces.detach().cpu().numpy(),
         )
         if self.compute_stress:
-            self.results.update(stress=stresses.detach().numpy() * self.stress_weight)
+            self.results.update(stress=stresses.detach().cpu().numpy() * self.stress_weight)
         if self.compute_hessian:
-            self.results.update(hessian=hessians.detach().numpy())
+            self.results.update(hessian=hessians.detach().cpu().numpy())
 
 
 class Relaxer:

diff --git a/matgl/graph/compute.py b/matgl/graph/compute.py
@@ -22,7 +22,7 @@ def compute_3body(g: dgl.DGLGraph):
         n_triple_s (np.ndarray): number of three-body angles for each structure
     """
     n_atoms = g.num_nodes()
-    first_col = g.edges()[0].numpy().reshape(-1, 1)
+    first_col = g.edges()[0].cpu().numpy().reshape(-1, 1)
     all_indices = np.arange(n_atoms).reshape(1, -1)
     n_bond_per_atom = np.count_nonzero(first_col == all_indices, axis=0)
     n_triple_i = n_bond_per_atom * (n_bond_per_atom - 1)