Is this a new feature, an improvement, or a change to existing functionality?
Improvement
How would you describe the priority of this feature request
Medium
Please provide a clear description of problem you would like to solve.
It can't be trained a GraphCast model with higher batchsize than 1. It is not explained why and the functions ask for a tensor with shape [N, C, H, W].
Is this a new feature, an improvement, or a change to existing functionality?
Improvement
How would you describe the priority of this feature request
Medium
Please provide a clear description of problem you would like to solve.
It can't be trained a GraphCast model with higher batchsize than 1. It is not explained why and the functions ask for a tensor with shape [N, C, H, W].
Describe any alternatives you have considered
Allow to train higher batchsizes