JAX problems

[agbruno-git]

I am running the model using my own data read from ERA5, the dataset I am using to feed the model looks identical to the one stored on the google bucked.
The worked fine when I run it locally with the data used in the demo, but I am having problems with JAX when I run the model using my data.

I had this error message

TypeError: Argument 'dask.array<getitem, shape=(1, 2, 37, 721, 1440), dtype=float32, chunksize=(1, 2, 37, 721, 1440), chunktype=numpy.ndarray>' of type <class 'dask.array.core.Array'> is not a valid JAX type.

Anyone had the same problem?

[alvarosg]

I think this may be fixed by calling ".compute()" on all of the data (inputs, targets and forcings), before passing it to the model.

[Zappandy]

@alvarosg I can confirm this was fixed by calling compute before passing it to the predictor object. Feel free to close the issue.