text: Title: Question about code of `vit_for_small_dataset.py`
  Body: Hi, 
  
  Thanks for the outstanding efforts to bring ViT to the framework of PyTorch, which is meaningful for me to learn about it!
  
  However, when I review the code in `vit_for_small_dataset.py`, in the `SPT`, I find code in the line `86` to be weird.
  
  ```python
  patch_dim = patch_size * patch_size * 5 * channels
  ```
  
  I fail to understand the meaning of `5` here. Is it a typo? Based on the context here, it seems that `patch_size * patch_size * channels` is more appropriate.
  
  Thanks for your attention!