Skip to content

GitLab

Unverified Commit 5761ceb3 authored by Yih-Dar's avatar Yih-Dar Committed by GitHub
Browse files

Fix device issue in `UperNetModelIntegrationTest` (#21192) 



fix device
Co-authored-by: default avatarydshieh <ydshieh@users.noreply.github.com>
parent 35920c97
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
tests/models/upernet/test_modeling_upernet.py
View file @ 5761ceb3
......@@ -272,7 +272,7 @@ def prepare_img():
class UperNetModelIntegrationTest(unittest.TestCase):
def test_inference_swin_backbone(self):
processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-swin-tiny")
model = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-swin-tiny")
model = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-swin-tiny").to(torch_device)
image = prepare_img()
inputs = processor(images=image, return_tensors="pt").to(torch_device)
......@@ -290,7 +290,7 @@ class UperNetModelIntegrationTest(unittest.TestCase):
def test_inference_convnext_backbone(self):
processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-tiny")
model = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-tiny")
model = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-tiny").to(torch_device)
image = prepare_img()
inputs = processor(images=image, return_tensors="pt").to(torch_device)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment