From 95ba90157fbbee293e8d10ac108ced2d9b990cbc Mon Sep 17 00:00:00 2001
From: Nguyen Nhat Hoang <55668910+nhathoang0110@users.noreply.github.com>
Date: Sat, 23 Sep 2023 11:21:44 +0700
Subject: [PATCH] Update tiny_vit.py to fix bug

---
 timm/models/tiny_vit.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/timm/models/tiny_vit.py b/timm/models/tiny_vit.py
index 4b583658..96a88db7 100644
--- a/timm/models/tiny_vit.py
+++ b/timm/models/tiny_vit.py
@@ -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: