From 5761ceb35a9ae0bd9e49a59438c725c61cda4f10 Mon Sep 17 00:00:00 2001 From: Yih-Dar <2521628+ydshieh@users.noreply.github.com> Date: Thu, 19 Jan 2023 14:26:14 +0100 Subject: [PATCH] Fix device issue in `UperNetModelIntegrationTest` (#21192) fix device Co-authored-by: ydshieh --- tests/models/upernet/test_modeling_upernet.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/models/upernet/test_modeling_upernet.py b/tests/models/upernet/test_modeling_upernet.py index 2e9daec542..8845edd2d4 100644 --- a/tests/models/upernet/test_modeling_upernet.py +++ b/tests/models/upernet/test_modeling_upernet.py @@ -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)