Fix index tensor device mismatch by moving to CPU

[recisic]

Description

This PR ensures the index tensor is sent to the CPU for full unrolling SCF mode during batched computations on a CUDA device. The tensor conv = self.mixer.converged may be a CUDA tensor as it is derived from deltas in the mixer class:

dxtb/src/dxtb/_src/scf/unrolling/default.py

Lines 173 to 184 in 92a7e77

	conv = self.mixer.converged
	if conv.any():
	# Simultaneous convergence does not require culling.
	# Occurs if batch size equals amount of True in `conv`.
	if guess.shape[0] == conv.count_nonzero():
	q_converged = q
	converged[:] = True
	culled = False
	break
	# save all necessary variables for converged system
	iconv = idxs[conv]

dxtb/src/dxtb/_src/scf/mixer/base.py

Lines 146 to 166 in 92a7e77

	@property
	def converged(self) -> Tensor:
	"""
	Tensor of bools indicating convergence status of the system(s).
	A system is considered to have converged if the maximum absolute
	difference between the current and previous systems is less than
	the ``tolerance`` value.
	"""
	# Check that mixing has been conducted
	if self.delta is None:
	raise RuntimeError("Nothing has been mixed")
	if self._batch_mode == 0:
	delta_norm = torch.norm(self.delta)
	else:
	# norm goes over all dims except first (batch dimension)
	dims = tuple(range(-(self.delta.ndim - 1), 0))
	delta_norm = torch.norm(self.delta, dim=dims)
	return delta_norm < self.options["x_tol"]

Fixes #194.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 84.04%. Comparing base (92a7e77) to head (43fabc2).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #195   +/-   ##
=======================================
  Coverage   84.04%   84.04%           
=======================================
  Files         200      200           
  Lines        9846     9846           
  Branches     1125     1125           
=======================================
  Hits         8275     8275           
  Misses       1216     1216           
  Partials      355      355           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[marvinfriede]

Thanks for reporting and fixing the issue!

However, instead of moving to the CPU, I prefer to initialize everything on the correct device. Usually, all tensors receive the device keyword argument, but in the full SCF, there some tensors without. I found more instances in the full SCF code, which is why I created a separate PR (#196).