Update README.md

README.md

@@ -47,6 +47,8 @@ It is adapted for downscaling of **2-channel ERA5 data** (e.g., wind u and v com
 from super_image import EdsrModel, EdsrConfig
 from huggingface_hub import hf_hub_download
 import torch
+import xarray as xr
+import numpy as np
 
 # load config 
 config, _ = EdsrConfig.from_pretrained("lschmidt/edsr-dsc")
@@ -62,18 +64,18 @@ state_dict = torch.load(state_dict_path, map_location="cpu")
 model.load_state_dict(state_dict, strict=False)
 
 # create random input: must be a 4D tensor (B, C=2, H, W)
-inputs = torch.randn(1, 2, 64, 64)  # replace with coarse wind velocity fields
+inputs = torch.randn(1, 2, 40, 40)  # replace with coarse wind velocity fields
 
 # or use sample data
-import xarray as xr
-import numpy as np
-
-data_path = hf_hub_download(repo_id="lschmidt/edsr-dsc/test_data", filename="test_wind_velocities.nc")
+data_path = hf_hub_download(
+    repo_id="lschmidt/edsr-dsc",
+    filename="test_wind_velocities.nc",
+    subfolder="test_data"  
+)
 ds = xr.open_dataset(data_path)
 u = ds["u100"].values[0]
 v = ds["v100"].values[0]
-inputs = torch.from_numpy(np.stack([u, v], axis=0)).unsqueeze(0).float()  # shape (1, 2, H, W)
+inputs = torch.from_numpy(np.stack([u, v], axis=0)).unsqueeze(0).float()  
 
 # prediction
 outputs = model(inputs)
-