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: