Update distill.py to include device agnostic code for `distill_mlp` head and `distillation_token`#324

vivekh2000 wants to merge 1 commit intolucidrains:mainlucidrains/vit-pytorch:mainfrom

vivekh2000:patch-3vivekh2000/vit-pytorch:patch-3Copy head branch name to clipboard