From af4ba92fdd6f4a36e9fc99aedae64e82deab25d3 Mon Sep 17 00:00:00 2001 From: ash Date: Wed, 19 Mar 2025 17:30:25 +0800 Subject: [PATCH] fix set_device twice bug --- mindocr/models/backbones/yolov8_backbone.py | 3 ++- mindocr/models/heads/yolov8_head.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/mindocr/models/backbones/yolov8_backbone.py b/mindocr/models/backbones/yolov8_backbone.py index ea9121777..96d1af2a4 100644 --- a/mindocr/models/backbones/yolov8_backbone.py +++ b/mindocr/models/backbones/yolov8_backbone.py @@ -1,3 +1,4 @@ +import os import numpy as np import mindspore as ms @@ -119,7 +120,7 @@ def yolov8_backbone( def test_yolo_backbone(): ms.set_context(mode=ms.PYNATIVE_MODE) - ms.set_device("Ascend", 3) + ms.set_device("Ascend", os.environ.get("DEVICE_ID", 0)) ms.set_seed(0) network = YOLOv8Backbone() diff --git a/mindocr/models/heads/yolov8_head.py b/mindocr/models/heads/yolov8_head.py index 0d3a81edd..c4b3c8ed3 100644 --- a/mindocr/models/heads/yolov8_head.py +++ b/mindocr/models/heads/yolov8_head.py @@ -1,4 +1,5 @@ import math +import os import numpy as np @@ -125,7 +126,7 @@ def yolov8_head(nc=5, reg_max=16, stride=None, in_channels=None) -> YOLOv8Head: def test_yolov8_head(): ms.set_context(mode=ms.PYNATIVE_MODE) - ms.set_device("Ascend", 3) + ms.set_device("Ascend", os.environ.get("DEVICE_ID", 0)) ms.set_seed(0) network = yolov8_head()