Update augvit_model.py

augvit_model.py

@@ -117,7 +117,7 @@ class AUGViT(Model):
         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',trainable=True)
+        self.pos_embedding = tf.Variable(initial_value=tf.random.normal([1, num_patches + 1, dim]),name='pos_emb',trainable=True)
         self.cls_token = tf.Variable(initial_value=tf.random.normal([1, 1, dim]),name='cls',trainable=True)
         self.dropout = nn.Dropout(rate=emb_dropout,name='drop')
 
@@ -141,8 +141,10 @@ class AUGViT(Model):
         )
         x = tf.concat([cls_tokens, x], axis=1)
         print(x.shape,cls_tokens.shape )
-        # x += self.pos_embedding[:, :(n+1 )]
-        print(x.shape)
+        pos= self.pos_embedding[:, :(n+1 )]
+        x += pos
+        
+        print(x.shape,pos.shape,self.pos_embedding.shape)
         x = self.dropout(x, training=training)
         print(x.shape)