Update README.md

README.md

@@ -34,11 +34,11 @@ data_config = timm.data.resolve_model_data_config(model)
 transforms = timm.data.create_transform(**data_config, is_training=False)
 
 with torch.no_grad():
-    output = model(transforms(img).unsqueeze(0))[0].softmax(dim=-1).cpu()
+    output = model(transforms(img).unsqueeze(0)).softmax(dim=-1).cpu()
 
 class_names = model.pretrained_cfg["label_names"]
-print("Probabilities:", output)
-print("Class:", class_names[output.argmax()])
+print("Probabilities:", output[0])
+print("Class:", class_names[output[0].argmax()])
 ```