Update README.md

README.md

@@ -41,10 +41,13 @@ datasets:
 <!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
 
 ```python
+from typing import Dict
+
 import numpy as np
 from datasets import load_dataset
 from matplotlib import cm
 from PIL import Image
+from torch import Tensor
 from transformers import AutoImageProcessor, AutoModel
 
 model = AutoModel.from_pretrained("RGBD-SOD/bbsnet", trust_remote_code=True)
@@ -82,7 +85,7 @@ preprocessed_sample = {
     'depth': tensor([[[[0.9529, 0....3490]]]])
 }
 """
-preprocessed_sample = image_processor.preprocess(sample)
+preprocessed_sample: Dict[str, Tensor] = image_processor.preprocess(sample)
 
 """
 2. Prediction step
@@ -91,7 +94,9 @@ output = {
     'logits': tensor([[[[-5.1966, ...ackward0>)
 }
 """
-output = model(preprocessed_sample["rgb"], preprocessed_sample["depth"])
+output: Dict[str, Tensor] = model(
+    preprocessed_sample["rgb"], preprocessed_sample["depth"]
+)
 
 """
 3. Postprocessing step