Gencast - num_steps_per_chunk > 1 breaks rollout.chunked_prediction_generator_multiple_runs

Running examples with num_steps_per_chunk = 2 results in the following error in example notebooks with the 1p0 model:
```
ValueError: 'grid2mesh_gnn/~_networks_builder/encoder_nodes_grid_nodes_mlp/~/linear_0/w' with retrieved shape (267, 512) does not match shape=[355, 512] dtype=dtype('float32')
```
Is the step chunking working?

Thanks!