allow channels to be customizable for cvt

Co-authored-by: Artem Lukin <artyom.lukin98@gmail.com>

README.md

@@ -93,7 +93,7 @@ preds = v(img) # (1, 1000)
 - `image_size`: int.  
 Image size. If you have rectangular images, make sure your image size is the maximum of the width and height
 - `patch_size`: int.  
-Number of patches. `image_size` must be divisible by `patch_size`.  
+Size of patches. `image_size` must be divisible by `patch_size`.  
 The number of patches is: ` n = (image_size // patch_size) ** 2` and `n` **must be greater than 16**.
 - `num_classes`: int.  
 Number of classes to classify.
@@ -777,7 +777,7 @@ pred = mbvit_xs(img) # (1, 1000)
 
 <img src="./images/xcit.png" width="400px"></img>
 
-This <a href="https://arxiv.org/abs/2106.09681">paper</a> introduces the cross correlation attention (abbreviated XCA). One can think of it as doing attention across the features dimension rather than the spatial one (another perspective would be a dynamic 1x1 convolution, the kernel being attention map defined by spatial correlations).
+This <a href="https://arxiv.org/abs/2106.09681">paper</a> introduces the cross covariance attention (abbreviated XCA). One can think of it as doing attention across the features dimension rather than the spatial one (another perspective would be a dynamic 1x1 convolution, the kernel being attention map defined by spatial correlations).
 
 Technically, this amounts to simply transposing the query, key, values before executing cosine similarity attention with learned temperature.
 

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'vit-pytorch',
   packages = find_packages(exclude=['examples']),
-  version = '1.6.0',
+  version = '1.6.3',
   license='MIT',
   description = 'Vision Transformer (ViT) - Pytorch',
   long_description_content_type = 'text/markdown',

vit_pytorch/__init__.py

@@ -1,10 +1,3 @@
-import torch
-from packaging import version
-
-if version.parse(torch.__version__) >= version.parse('2.0.0'):
-    from einops._torch_specific import allow_ops_in_compiled_graph
-    allow_ops_in_compiled_graph()
-
 from vit_pytorch.vit import ViT
 from vit_pytorch.simple_vit import SimpleViT
 

vit_pytorch/cvt.py

@@ -140,12 +140,13 @@ class CvT(nn.Module):
         s3_heads = 6,
         s3_depth = 10,
         s3_mlp_mult = 4,
-        dropout = 0.
+        dropout = 0.,
+        channels = 3
     ):
         super().__init__()
         kwargs = dict(locals())
 
-        dim = 3
+        dim = channels
         layers = []
 
         for prefix in ('s1', 's2', 's3'):

vit_pytorch/vit_1d.py

@@ -10,7 +10,7 @@ class FeedForward(nn.Module):
     def __init__(self, dim, hidden_dim, dropout = 0.):
         super().__init__()
         self.net = nn.Sequential(
-            nn.Layernorm(dim),
+            nn.LayerNorm(dim),
             nn.Linear(dim, hidden_dim),
             nn.GELU(),
             nn.Dropout(dropout),