Update src/transformers/models/textnet/modeling_textnet.py

Co-authored-by: Pavel Iakubovskii <[email protected]>
huggingface · Dec 21, 2024 · ed36fef · ed36fef
1 parent 5f29454
commit ed36fef
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/transformers/models/textnet/modeling_textnet.py b/src/transformers/models/textnet/modeling_textnet.py
@@ -346,9 +346,10 @@ def forward(
 
         Examples:
         ```python
-        >>> from transformers import TextNetForImageClassification,TextNetImageProcessor
-        >>> from PIL import Image
+        >>> import torch
         >>> import requests
+        >>> from transformers import TextNetForImageClassification, TextNetImageProcessor
+        >>> from PIL import Image
 
         >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
         >>> image = Image.open(requests.get(url, stream=True).raw)
@@ -357,7 +358,8 @@ def forward(
         >>> model = TextNetForImageClassification.from_pretrained("Raghavan/textnet-base")
 
         >>> inputs = processor(images=image, return_tensors="pt", size={"height": 640, "width": 640})
-        >>> outputs = model(**inputs)
+        >>> with torch.no_grad():
+        ...     outputs = model(**inputs)
         >>> outputs.logits.shape
         torch.Size([1, 2])
         ```"""