make nest resilient to dimension that are not divisible by number of heads

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'vit-pytorch',
   packages = find_packages(exclude=['examples']),
-  version = '0.19.0',
+  version = '0.19.1',
   license='MIT',
   description = 'Vision Transformer (ViT) - Pytorch',
   author = 'Phil Wang',

vit_pytorch/nest.py

@@ -48,16 +48,16 @@ class FeedForward(nn.Module):
 class Attention(nn.Module):
     def __init__(self, dim, heads = 8, dropout = 0.):
         super().__init__()
-        assert (dim % heads) == 0, 'dimension must be divisible by number of heads'
         dim_head = dim // heads
+        inner_dim = dim_head * heads
         self.heads = heads
         self.scale = dim_head ** -0.5
 
         self.attend = nn.Softmax(dim = -1)
-        self.to_qkv = nn.Conv2d(dim, dim * 3, 1, bias = False)
+        self.to_qkv = nn.Conv2d(dim, inner_dim * 3, 1, bias = False)
 
         self.to_out = nn.Sequential(
-            nn.Conv2d(dim, dim, 1),
+            nn.Conv2d(inner_dim, dim, 1),
             nn.Dropout(dropout)
         )