Update README.md

README.md

@@ -70,7 +70,7 @@ for sample in tqdm(data):
     ).to("cuda")
 
     with torch.no_grad():
-        generated_label = model(**inputs).logits.argmax()
+        generated_label = model(**inputs).logits.argmax().item()
 
     if generated_label == true_label:
         passes += 1