TypeError: __init__() got an unexpected keyword argument 'qk_scale'

[liuchanglab]

We run the code to load the model and get the following error:
File "/mlx_devbox/users/liuchang.lab/repo/22056/mae/inference.py", line 96, in
model_mae = prepare_model(chkpt_dir, 'mae_vit_large_patch16')
File "/mlx_devbox/users/liuchang.lab/repo/22056/mae/inference.py", line 29, in prepare_model
model = getattr(models_mae, arch)()
File "/mlx_devbox/users/liuchang.lab/repo/22056/mae/models_mae.py", line 232, in mae_vit_large_patch16_dec512d8b
model = MaskedAutoencoderViT(
File "/mlx_devbox/users/liuchang.lab/repo/22056/mae/models_mae.py", line 39, in init
self.blocks = nn.ModuleList([
File "/mlx_devbox/users/liuchang.lab/repo/22056/mae/models_mae.py", line 40, in
Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer)

How to fix it?