Feature: Easy switch between GPU/CPU for forcefields

[JaGeo]

For workstations, switching between GPUs and CPUs should be possible for the force field workflows. Do we currently have an easy option within the workflow to do so?

We typically run these jobs on either GPUs or CPUs and rarely have an issue. Did anyone of you run into this? @esoteric-ephemera @janosh ? Especially in the context of CHGNet?

[janosh]

Are you asking about switching device mid-workflow? I've never tried or needed that, so can't comment.

To initialize a workflow on different device, there's the use_device kwarg.

I was going to add a new parameter model_kwargs to the ForceFieldMakers be able to customize model loading including use_device

[JaGeo]

Are you asking about switching device mid-workflow? I've never tried or needed that, so can't comment.

To initialize a workflow on different device, there's the use_device kwarg.

I was going to add a new parameter model_kwargs to the ForceFieldMakers be able to customize model loading including use_device

I am asking about initialization. Would be great to have that parameter

[esoteric-ephemera]

@janosh you can already pass that switch through calculator_kwargs if it's a kwarg of the ML model's calculator

[janosh]

sweet, i'd forgotten you added that in #722.

looks like there are some doc string vestiges from a MACE-specific version of that which I'd planned to standardize across models at some point. i'll remove those

atomate2/src/atomate2/forcefields/jobs.py

Lines 406 to 409 in 3486792

	model_kwargs: dict[str, Any]
	Further keywords (e.g. device, default_dtype, model) for
	:obj:`mace.calculators.MACECalculator()'`.
	"""

[JaGeo]

@esoteric-ephemera @janosh thank you both! I will then close this as resolved!

[esoteric-ephemera]

Oh yeah if you don't mind doing that, that would be great. I also forgot to copy the base doc str over from the parent Relax, Static, and MD forcefield makers. If you don't mind adding that too, would be super appreciated