Update tiny_vit.py to fix bug

This commit is contained in:
Nguyen Nhat Hoang 2023-09-23 11:21:44 +07:00 committed by Ross Wightman
parent a9c29a05de
commit 95ba90157f

View File

@ -378,6 +378,7 @@ class TinyVitStage(nn.Module):
super().__init__()
self.depth = depth
self.out_dim = out_dim
# patch merging layer
if downsample is not None: