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,7 +141,7 @@ class AUGViT(Model):
         # )
         # x = tf.concat([cls_tokens, x], axis=1)
         # print(x.shape,cls_tokens.shape )
-        # x += self.pos_embedding[:, :(n + 1)]
+        x += self.pos_embedding[:, :(n + 1)]
         print(x.shape)
         x = self.dropout(x, training=training)
         print(x.shape)