From 43fabc2c20e8dacd0df5b5a532b999a657f61b20 Mon Sep 17 00:00:00 2001 From: Juno Nam Date: Mon, 13 Jan 2025 14:44:55 -0500 Subject: [PATCH] Fix index tensor device mismatch by moving to CPU --- src/dxtb/_src/scf/unrolling/default.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/dxtb/_src/scf/unrolling/default.py b/src/dxtb/_src/scf/unrolling/default.py index 89672327..49be80dc 100644 --- a/src/dxtb/_src/scf/unrolling/default.py +++ b/src/dxtb/_src/scf/unrolling/default.py @@ -181,7 +181,7 @@ def scf( break # save all necessary variables for converged system - iconv = idxs[conv] + iconv = idxs[conv.cpu()] q_converged[iconv, :mpdim, :norb] = q[conv, ..., :] ch[iconv, :norb, :norb] = self._data.hamiltonian[conv, :, :] cevecs[iconv, :norb, :norb] = self._data.evecs[conv, :, :] @@ -247,7 +247,7 @@ def scf( # cull local variables q = q[~conv, :mpdim, :norb] - idxs = idxs[~conv] + idxs = idxs[~conv.cpu()] if self._data.charges["mono"] is not None: self._data.charges["mono"] = torch.Size(