From b5fc0c7c25ab89ce223a4252e38fbbe4bfc7fe54 Mon Sep 17 00:00:00 2001
From: Marvin Friede <51965259+marvinfriede@users.noreply.github.com>
Date: Thu, 6 Jun 2024 07:48:52 +0200
Subject: [PATCH] Move data to CUDA before indexing (#57)

---
 src/tad_dftd4/damping/atm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/tad_dftd4/damping/atm.py b/src/tad_dftd4/damping/atm.py
index 4e4365c..2db6901 100644
--- 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