Update augvit_model.py

augvit_model.py

@@ -117,8 +117,8 @@ 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.cls_token = tf.Variable(initial_value=tf.random.normal([1, 1, dim]),name='cls',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')
 
         self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout=dropout,name='trans')
@@ -135,13 +135,13 @@ class AUGViT(Model):
         x = self.patch_den(x)
         b, n, d = x.shape
         print(x.shape)
-        # cls_tokens = tf.cast(
-        #     tf.broadcast_to(self.cls_token, [b, 1, d]),
-        #     dtype=x.dtype,
-        # )
-        # x = tf.concat([cls_tokens, x], axis=1)
-        # print(x.shape,cls_tokens.shape )
-        x += self.pos_embedding[:, :(n )]
+        cls_tokens = tf.cast(
+            tf.broadcast_to(self.cls_token, [b, 1, d]),
+            dtype=x.dtype,
+        )
+        x = tf.concat([cls_tokens, x], axis=1)
+        print(x.shape,cls_tokens.shape )
+        # x += self.pos_embedding[:, :(n+1 )]
         print(x.shape)
         x = self.dropout(x, training=training)
         print(x.shape)