Predictions coming out NAN

I am trying to run the 0.25 degree model and make predictions for 48 hours ahead.

I downloaded the weights and normalization

This is my input

Targets

Forcings

I am trying to make predictions like this:

`with open(r'/model/params/params_GraphCast - ERA5 1979-2017 - resolution 0.25 - pressure levels 37 - mesh 2to6 - precipitation input and output.npz', 'rb') as model:
ckpt = checkpoint.load(model, graphcast.CheckPoint)

params = ckpt.params
state = {}
model_config = ckpt.model_config
task_config = ckpt.task_config

with open(r'/model/stats/stats_diffs_stddev_by_level.nc', 'rb') as f:
diffs_stddev_by_level = xarray.load_dataset(f).compute()

with open(r'/model/stats/stats_mean_by_level.nc', 'rb') as f:
mean_by_level = xarray.load_dataset(f).compute()

with open(r'/model/stats/stats_stddev_by_level.nc', 'rb') as f:
stddev_by_level = xarray.load_dataset(f).compute()

def construct_wrapped_graphcast(model_config:graphcast.ModelConfig, task_config:graphcast.TaskConfig):
predictor = graphcast.GraphCast(model_config, task_config)
#predictor = casting.Bfloat16Cast(predictor)
predictor = normalization.InputsAndResiduals(predictor, diffs_stddev_by_level = diffs_stddev_by_level, mean_by_level = mean_by_level, stddev_by_level = stddev_by_level)
predictor = autoregressive.Predictor(predictor, gradient_checkpointing = True)
return predictor

@hk.transform_with_state
def run_forward(model_config, task_config, inputs, targets_template, forcings):
predictor = construct_wrapped_graphcast(model_config, task_config)
return predictor(inputs, targets_template = targets_template, forcings = forcings)

def with_configs(fn):
return functools.partial(fn, model_config = model_config, task_config = task_config)

def with_params(fn):
return functools.partial(fn, params = params, state = state)

def drop_state(fn):
return lambda **kw: fn(**kw)[0]

run_forward_jitted = drop_state(with_params(jax.jit(with_configs(run_forward.apply))))

class Predictor:

@classmethod
def predict(cls, inputs, targets, forcings) -> xarray.Dataset:
predictions = rollout.chunked_prediction(run_forward_jitted, rng = jax.random.PRNGKey(0), inputs = inputs, targets_template = targets, forcings = forcings)
return predictions`

At first I get this error: dtype argument must be a NumPy dtype, but it is a .

So then I comment out (#predictor = casting.Bfloat16Cast(predictor))

The the model runs and I get the predictions but all the values are NAN

I am using CPU only but I have lots of memory available and I was able to make a prediction for the same dates but on 1 degree data still using the 0.25 degree model's checkpoint.

If anyone knows why the model is outputting all NAN I would really appreciate the help.

Start with the checkpoint and normalization loading, then inspect construct_wrapped_graphcast and rollout.chunked_prediction while comparing the 0.25-degree inputs with the working 1-degree case. Trace where NaN values first appear after removing casting.Bfloat16Cast; done means the 0.25-degree prediction completes with finite values.

Written by the indexing model from the issue text.