Update vit.py to include LayerNorm in the MLP head which is missing

[vivekh2000]

I have included the LayerNorm layer in the classification head to identify it with the original implementation. It will be easy to compare now. Also, an isolated LayerNorm was included, which I have removed and inserted in the MLP head, again to make it more intuitive. For reference, you have included LayerNorm in the self.ditill_mlp attribute of the class DistillWrapper in distill.py. I hope these changes are acceptable.

Thank you for the excellent code base.