vit_patch: Patchification of input image not correctly checked

[Anuj07-dev]

Incorrect patch for input x unchecked in testcases:

    def forward(self, x):
      B, C, H, W = x.shape
      x_new = x.view(B, self.num_patches, -1)

      patch_embedding = self.proj(x_new)
      return patch_embedding

Correct patch for x:

    def forward(self, x):
      B, C, H, W = x.shape
      x_new = x.reshape(B, C, H//self.patch_size, self.patch_size, W//self.patch_size, self.patch_size).permute(0, 2, 4, 1, 3, 5).reshape(B, self.num_patches, C * self.patch_size * self.patch_size)

      patch_embedding = self.proj(x_new)
      return patch_embedding