cleanup

Author: lucidrains

vit_pytorch/vit_with_keel_post_ln.py

@@ -189,7 +189,7 @@ def forward(self, img):
 
         x = self.transformer(x)
 
-        if self.mlp_head is None:
+        if not exists(self.mlp_head):
             return x
 
         x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]