Visualize

Hello,

I am using this command, but it does not save any masked images.
python3 main.py --model deit_small_patch16_shrink_base --fuse_token --base_keep_rate 0.7 --visualize_mask --n_visualization 64 --resume checkpoint --data-path /path/to/imagenet

I think this part is not working in the code:
if args.visualize_mask:
visualize_mask(data_loader_val, model, device, args.output_dir, args.n_visualization, args.fuse_token)
return