Run graphcast on the correct device

[nbren12]

Earth-2 MIP Pull Request

Description

When multiple GPUs are present graphcast was not running on the specified device. Jax places the model on the device of the input data it is passed. However, when converting the torch tensor to an xarray.Dataset of jax arrays the data were being cast to CPU numpy arrays. Using xarray_jax.Variable resolves this.

Checklist

• [ ] I am familiar with the Contributing Guidelines.
• [ ] New or existing tests cover these changes.
• [ ] The documentation is up to date with these changes.
• [ ] The CHANGELOG.md is up to date with these changes.
• [ ] An issue is linked to this pull request.

Dependencies

[nbren12]

/blossom-ci

[nbren12]

I just verified this works by running the time_collection script on a selene node with 8 gpus and running nvidia-smi.

[nbren12]

/blossom-ci