shape mismatch between pretrained model and code

[maxmars1]

/home/myfolder/anaconda3/envs/OSEDiff/lib/python3.10/site-packages/timm/models/layers/init.py:48: FutureWarning: Importing from timm.models.layers is deprecated, please import via timm.layers
warnings.warn(f"Importing from {name} is deprecated, please import via timm.layers", FutureWarning)
/home/myfolder/anaconda3/envs/OSEDiff/lib/python3.10/site-packages/timm/models/hub.py:4: FutureWarning: Importing from timm.models.hub is deprecated, please import via timm.models
warnings.warn(f"Importing from {name} is deprecated, please import via timm.models", FutureWarning)
/home/myfolder/anaconda3/envs/OSEDiff/lib/python3.10/site-packages/timm/models/registry.py:4: FutureWarning: Importing from timm.models.registry is deprecated, please import via timm.models
warnings.warn(f"Importing from {name} is deprecated, please import via timm.models", FutureWarning)
/home/myfolder/anaconda3/envs/OSEDiff/lib/python3.10/site-packages/timm/models/helpers.py:7: FutureWarning: Importing from timm.models.helpers is deprecated, please import via timm.models
warnings.warn(f"Importing from {name} is deprecated, please import via timm.models", FutureWarning)
Traceback (most recent call last):
File "/home/myfolder/2025/OSEDiff/test_osediff.py", line 64, in
model = OSEDiff_test(args)
File "/home/myfolder/2025/OSEDiff/osediff.py", line 267, in init
self.load_ckpt(osediff)
File "/home/myfolder/2025/OSEDiff/osediff.py", line 294, in load_ckpt
p.data.copy_(model["state_dict_unet"][n])
RuntimeError: output with shape [4, 320, 1, 1] doesn't match the broadcast shape [4, 320, 4, 320]

I uses the latest code and pretrained model in README.md. any help?