From 2a4c8f4b65d5270419e539d2d3917f91075128c6 Mon Sep 17 00:00:00 2001
From: wangguanzhong <jerrywgz@126.com>
Date: Thu, 9 Jun 2022 17:28:27 +0800
Subject: [PATCH] keep device in export (#6159)

---
 ppdet/engine/export_utils.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/ppdet/engine/export_utils.py b/ppdet/engine/export_utils.py
index e1cf64638..9608699f1 100644
--- a/ppdet/engine/export_utils.py
+++ b/ppdet/engine/export_utils.py
@@ -55,7 +55,9 @@ MOT_ARCH = ['DeepSORT', 'JDE', 'FairMOT']
 def _prune_input_spec(input_spec, program, targets):
     # try to prune static program to figure out pruned input spec
     # so we perform following operations in static mode
+    device = paddle.get_device()
     paddle.enable_static()
+    paddle.set_device(device)
     pruned_input_spec = [{}]
     program = program.clone()
     program = program._prune(targets=targets)
@@ -66,7 +68,7 @@ def _prune_input_spec(input_spec, program, targets):
             pruned_input_spec[0][name] = spec
         except Exception:
             pass
-    paddle.disable_static()
+    paddle.disable_static(place=device)
     return pruned_input_spec
 
 
-- 
GitLab