Earth-2 Model Intercomparison Project (MIP) is a python framework that enables climate researchers and scientists to inter-compare AI models for weather and climate.
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.
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
Dependencies