HarshRangwala:

Invalid shape (1, 3, 224, 224) for image data


That first dimension should be squeezed out as an image should have 3 dimensions: number of channels, height, and width. (i.e. (3, 224, 224)).
Try img = img.squeeze() before calling ax.imshow(img)