Forecasting with graphcast_small

I'm trying to make a forecast using graphcast_small and I'm getting a shape dismatch error. Next, using padding, I adjust the dimensions of the layers. But I encounter the following error:
"TypeError: scan body function carry input and carry output must have the same pytree structure, but they differ".

Please help me correct the errors. The code can be viewed at the link: https://github.com/Michael-Sharma/projects/blob/main/draft.ipynb