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