The problem is that when I run “python finetune.py --train && python finetune.py --prune” which is the forward pass, these two commands do not give me tensor size mismatch issues.
I will definitely look at predict.py again, but this file does not mess around with the model internally itself.
So, I guess the problems lie at prune.py 1 which messes the original SqueezeNet model.