Need Clarification regarding prune_deit() in - https://github.com/locuslab/wanda/blob/main/image_classifiers/main.py

[solomonmanuelraj]

Hi Team,

in the image classifier code, ( https://github.com/locuslab/wanda/blob/main/image_classifiers/main.py)

in line number 324 to 332,

############################################################################# tick = time.time() if args.sparsity != 0: with torch.no_grad(): if "convnext" in args.model: prune_convnext(args, model, calib_data, device) elif "vit" in args.model: prune_vit(args, model, calib_data, device) elif "deit" in args.model: prune_vit(args, model, calib_data, device) ########################################################################### in line number 329 and 331 conditions both the time you are calling prune_vit(). why we are not calling prune_deit() even though we have imported that in the from prune_utils import prune_convnext, prune_deit, prune_vit, check_sparsity line.

regards