Update README.md

README.md

@@ -61,13 +61,13 @@ MODEL_NAME = "p1atdev/wd-swinv2-tagger-v3-hf"
 model = AutoModelForImageClassification.from_pretrained(
     MODEL_NAME,
 )
-model.eval()
 processor = AutoImageProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)
 
 image = Image.open("sample.webp")
 inputs = processor.preprocess(image, return_tensors="pt")
 
-outputs = model(**inputs.to(model.device, model.dtype))
+with torch.no_grad():
+  outputs = model(**inputs.to(model.device, model.dtype))
 logits = torch.sigmoid(outputs.logits[0])
 
 # get probabilities