always fix head dimension at 64, so users do not trip up with imprope…

…r ratio of dimension to heads

examples/enwik8_simple/train.py

@@ -82,7 +82,7 @@ def __len__(self):
     model.train()
 
     for __ in range(GRADIENT_ACCUMULATE_EVERY):
-        loss = model(next(train_loader), return_loss = True)
+        loss = model(next(train_loader))
         loss.backward()
 
     print(f'training loss: {loss.item()}')
@@ -93,7 +93,7 @@ def __len__(self):
     if i % VALIDATE_EVERY == 0:
         model.eval()
         with torch.no_grad():
-            loss = model(next(val_loader), return_loss = True)
+            loss = model(next(val_loader))
             print(f'validation loss: {loss.item()}')
 
     if i % GENERATE_EVERY == 0:

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'x-transformers',
   packages = find_packages(exclude=['examples']),
-  version = '0.0.3',
+  version = '0.0.4',
   license='MIT',
   description = 'X-Transformers - Pytorch',
   author = 'Phil Wang',

x_transformers/x_transformers.py

@@ -49,17 +49,17 @@ def forward(self, x):
         return self.net(x)
 
 class Attention(nn.Module):
-    def __init__(self, dim, heads = 8, causal = False, mask = None):
+    def __init__(self, dim, dim_head = 64, heads = 8, causal = False, mask = None):
         super().__init__()
-        assert (dim % heads) == 0, 'dimension must be divisible by number of heads'
-        self.scale = (dim // heads) ** -0.5
+        self.scale = dim_head ** -0.5
         self.heads = heads
         self.causal = causal
         self.mask = mask
 
-        self.to_q = nn.Linear(dim, dim, bias = False)
-        self.to_kv = nn.Linear(dim, dim * 2, bias = False)
-        self.to_out = nn.Linear(dim, dim)
+        inner_dim = dim_head * heads
+        self.to_q = nn.Linear(dim, inner_dim, bias = False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False)
+        self.to_out = nn.Linear(inner_dim, dim)
 
     def forward(self, x, context = None, mask = None, context_mask = None):
         b, n, _, h, device = *x.shape, self.heads, x.device