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()