update example with vit with patch dropout

Author: lucidrains

README.md

@@ -17,7 +17,7 @@ $ pip install coca-pytorch
 First install the `vit-pytorch` for the image encoder, which needs to be pretrained
 
 ```bash
-$ pip install vit-pytorch>=0.38.1
+$ pip install vit-pytorch>=0.40.2
 ```
 
 Then
@@ -27,17 +27,18 @@ import torch
 
 # import vision transformer
 
-from vit_pytorch import SimpleViT
+from vit_pytorch.simple_vit_with_patch_dropout import SimpleViT
 from vit_pytorch.extractor import Extractor
 
-vit = ViT(
+vit = SimpleViT(
     image_size = 256,
     patch_size = 32,
     num_classes = 1000,
     dim = 1024,
     depth = 6,
     heads = 16,
-    mlp_dim = 2048
+    mlp_dim = 2048,
+    patch_dropout = 0.5  # https://arxiv.org/abs/2212.00794
 )
 
 vit = Extractor(vit, return_embeddings_only = True, detach = False)