Update augvit_model.py

augvit_model.py

@@ -134,11 +134,14 @@ class AUGViT(Model):
         x = self.patch_embedding(img)
         x = self.patch_den(x)
         b, n, d = x.shape
-
+        print(x.shape)
         cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=b)
         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)
 
         x = self.transformer(x, training=training)