Update README.md

README.md

@@ -6,8 +6,12 @@ license: mit
 
 ## VitMatte
 The demo files are [here](https://github.com/hustvl/ViTMatte/tree/main/demo)
+
+
+
 ```python
 import torch
+from torchvision.transforms import functional as F
 
 image = Image.open("./retriever_rgb.png").convert("RGB")
 image = F.to_tensor(image).unsqueeze(0).to("cuda").half()
@@ -22,4 +26,6 @@ alpha = model(input)
 
 output = F.to_pil_image(predictions)
 output.save("./predicted.png")
-```
+```
+This is the output:
+<img width=400 src="https://cdn-uploads.huggingface.co/production/uploads/62ff65702979d8fc339b0905/2LdOFka5RXucDCWzzVxwl.png"/>