May be a bug?

[RachelTeamo]

Hi～Thanks for your code.
But in this line it may be a bug. I try your code in my tasks, and the image shape is [382,382]. But args.token_size=16 means self.num_patches=256, but 382/16 * 382/16 = 512. In your demo, your image shape is [256,256] but 256/16 =16 is a coincide. When the image is not [256,256], the number of patches may not right.

So I guess maybe we should use args.input_size//args.token_size instead of args.token_size.

Please take a look, excuse me.