Move data to CUDA before indexing (#57)

dftd4 · Jun 6, 2024 · b5fc0c7 · b5fc0c7
1 parent 4f996e8
commit b5fc0c7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/tad_dftd4/damping/atm.py b/src/tad_dftd4/damping/atm.py
@@ -103,7 +103,7 @@ def get_atm_dispersion(
         torch.abs(c6.unsqueeze(-1) * c6.unsqueeze(-2) * c6.unsqueeze(-3)),
     )
 
-    rad = data.R4R2[numbers]
+    rad = data.R4R2.to(**dd)[numbers]
     radii = rad.unsqueeze(-1) * rad.unsqueeze(-2)
     temp = a1 * storch.sqrt(3.0 * radii) + a2