ZF/vit-pytorch
1
0
Fork 0
mirror of https://github.com/lucidrains/vit-pytorch.git synced 2025-12-30 08:02:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

cleanup stray print

Browse Source
This commit is contained in:
Phil Wang
2020-10-07 11:22:45 -07:00
parent f7123720c3
commit d66b29e4cf
2 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
setup.py
View File

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
setup(
name = 'vit-pytorch',
packages = find_packages(),
version = '0.0.4',
version = '0.0.5',
license='MIT',
description = 'Vision Transformer (ViT) - Pytorch',
author = 'Phil Wang',

1
vit_pytorch/vit_pytorch.py
View File

@@ -46,7 +46,6 @@ class Attention(nn.Module):
if mask is not None:
mask = F.pad(mask.flatten(1), (1, 0), value = True)
print(mask.shape[-1], dots.shape[-1])
assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions'
mask = mask[:, None, :] * mask[:, :, None]
dots.masked_fill_(~mask, float('-inf'))