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 + 1)]
+        # 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)