Update augvit_model.py

augvit_model.py

@@ -113,10 +113,9 @@ class AUGViT(Model):
         num_patches = (image_height // patch_height) * (image_width // patch_width)
         assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
 
-        self.patch_embedding = Sequential([
-            Rearrange('b (h p1) (w p2) c -> b (h w) (p1 p2 c)', p1=patch_height, p2=patch_width),
-            nn.Dense(units=dim,name='patchden')
-        ], name='patch_embedding')
+        self.patch_embedding = Rearrange('b (h p1) (w p2) c -> b (h w) (p1 p2 c)', p1=patch_height, p2=patch_width)
+        self.patch_den=    nn.Dense(units=dim,name='patchden')
+        
 
         self.pos_embedding = tf.Variable(initial_value=tf.random.normal([1, num_patches + 1, dim]),name='pos_emb')
         self.cls_token = tf.Variable(initial_value=tf.random.normal([1, 1, dim]),name='cls')
@@ -133,6 +132,7 @@ class AUGViT(Model):
 
     def call(self, img, training=True, **kwargs):
         x = self.patch_embedding(img)
+        x = self.patch_den(x)
         b, n, d = x.shape
 
         cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=b)