Dealing with pre-trained model for grayscale images
print(grayscale_batch.shape)  # (64, 224, 224)    

rgb_batch = np.repeat(grayscale_batch[..., np.newaxis], 3, -1)    
print(rgb_batch.shape)  # (64, 224, 224, 3)