Loading pertaining model fails on custom "load_state_dict_ function

[NimrodShabtay]

Hi, @MzeroMiko
Thanks for a super interesting work.
I am trying to use the pre-trained classification model vssm_base_0229_ckpt_epoch_237.pth

i am building the model using the config classification/configs/vssm/vmambav2_base_224.yaml

But when trying to load the model I get:

vmamba_model._load_from_state_dict(ckpt["model"])
TypeError: VSSM._load_from_state_dict() missing 6 required positional arguments: 'prefix', 'local_metadata', 'strict', 'missing_keys', 'unexpected_keys', and 'error_msgs'

How can I just load the model and run it?
Thanks