🚀[FEA]: Increment GraphCast BatchSize

[Javier-Jimenez99]

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 with higher batchsizes

[divijkhaitan]

@mnabian Is this feature still being developed?