Update example

README.md

@@ -60,6 +60,8 @@ pip install microsoft-aurora
 Run the pretrained small model on random data:
 
 ```python
+from datetime import datetime
+
 import torch
 
 from aurora import AuroraSmall, Batch, Metadata
@@ -70,7 +72,7 @@ model.load_checkpoint("wbruinsma/aurora", "aurora-0.25-small-pretrained.ckpt")
 
 batch = Batch(
     surf_vars={k: torch.randn(1, 2, 16, 32) for k in ("2t", "10u", "10v", "msl")},
-    static_vars={k: torch.randn(1, 2, 16, 32) for k in ("lsm", "z", "slt")},
+    static_vars={k: torch.randn(16, 32) for k in ("lsm", "z", "slt")},
     atmos_vars={k: torch.randn(1, 2, 4, 16, 32) for k in ("z", "u", "v", "t", "q")},
     metadata=Metadata(
         lat=torch.linspace(90, -90, 17)[:-1],  # Cut off the south pole.