diff --git a/setup.py b/setup.py index ca864bc..a815baa 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ from setuptools import setup, find_packages setup( name = 'vit-pytorch', packages = find_packages(exclude=['examples']), - version = '0.26.5', + version = '0.26.6', license='MIT', description = 'Vision Transformer (ViT) - Pytorch', author = 'Phil Wang', diff --git a/vit_pytorch/recorder.py b/vit_pytorch/recorder.py index 3d3ad7e..7ecac2c 100644 --- a/vit_pytorch/recorder.py +++ b/vit_pytorch/recorder.py @@ -55,5 +55,5 @@ class Recorder(nn.Module): target_device = self.device if self.device is not None else img.device recordings = tuple(map(lambda t: t.to(target_device), self.recordings)) - attns = torch.stack(recordings, dim = 1) + attns = torch.stack(recordings, dim = 1) if len(recordings) > 0 else None return pred, attns