diff --git a/README.md b/README.md
index e4feb996e..c5bf501b1 100644
--- a/README.md
+++ b/README.md
@@ -112,8 +112,8 @@ with torch.no_grad():
     text_features = model.encode_text(text_inputs)
 
 # Pick the top 5 most similar labels for the image
-image_features /= image_features.norm(dim=-1, keepdim=True)
-text_features /= text_features.norm(dim=-1, keepdim=True)
+image_features /= image_features.norm(dim=1, keepdim=True)
+text_features /= text_features.norm(dim=1, keepdim=True)
 similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1)
 values, indices = similarity[0].topk(5)