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

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?