Use xarray.Dataset/Array.sizes instead of .dims

[timothyas]

There are a few places where the syntax xarray.Dataset.dims[varname] is used to get the length of e.g. varname. This produces the following warning of to-be-deprecated behavior

graphcast/graphcast/autoregressive.py:115: FutureWarning: The return type of `Dataset.dims` will be changed to return a set of dimension names in future, in order to be more consistent with `DataArray.dims`. To access a mapping from dimension names to lengths, please use `Dataset.sizes`.

(for reference this is with xarray version 2024.5.0)

This PR makes the replacement suggested by the warning. I think these are all the spots where this functionality of .dims is used, it's at least a good start.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.