NVIDIA / earth2mip

Earth-2 Model Intercomparison Project (MIP) is a python framework that enables climate researchers and scientists to inter-compare AI models for weather and climate.
https://nvidia.github.io/earth2mip/
Apache License 2.0
187 stars 41 forks source link

🐛[BUG]: README example not working with e2mip://fcn #62

Closed nbren12 closed 11 months ago

nbren12 commented 11 months ago

Version

main

On which installation method(s) does this occur?

Pip

Describe the issue

I get this bug

In [3]: >>> import datetime
   ...: >>> from earth2mip.networks import get_model
   ...: >>> from earth2mip.initial_conditions import cds
   ...: >>> from earth2mip.inference_ensemble import run_basic_inference
   ...: >>> time_loop  = get_model("e2mip://fcn", device="cuda:0")
   ...: >>> data_source = cds.DataSource(time_loop.in_channel_names)
   ...: >>> ds = run_basic_inference(time_loop, n=10, data_source=data_source, time=datetime.datetime(2018, 1, 1))
   ...: >>> ds.chunk()

File ~/.local/lib/python3.10/site-packages/modulus/models/afno/afno.py:372, in PatchEmbed.forward(self, x)
    369 def forward(self, x: Tensor) -> Tensor:
    370     B, C, H, W = x.shape
    371     assert (
--> 372         H == self.img_size[0] and W == self.img_size[1]
    373     ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
    374     x = self.proj(x).flatten(2).transpose(1, 2)
    375     return x

AssertionError: Input image size (721*1440) doesn't match model (720*1440).

this is probably a bug in run_basic_inference.

Environment details

No response