refine set_device (#4102)

e4982704 · wangguanzhong · GitHub · 69e909b4 · e4982704
隐藏空白更改
内联并排

Showing with 3 addition and 2 deletion

ppdet/modeling/proposal_generator/target.py ppdet/modeling/proposal_generator/target.py +3 -2

未找到文件。
--- a/ppdet/modeling/proposal_generator/target.py
+++ b/ppdet/modeling/proposal_generator/target.py
@@ -74,8 +74,9 @@ def label_box(anchors,
              is_crowd=None,
              assign_on_cpu=False):
    if assign_on_cpu:
-        with paddle.fluid.framework._dygraph_place_guard(paddle.CPUPlace()):
-            iou = bbox_overlaps(gt_boxes, anchors)
+        paddle.set_device("cpu")
+        iou = bbox_overlaps(gt_boxes, anchors)
+        paddle.set_device("gpu")
    else:
        iou = bbox_overlaps(gt_boxes, anchors)
    n_gt = gt_boxes.shape[0]