From 39ee8eb3daad8477e38eb1f8ebb621d1be090df6 Mon Sep 17 00:00:00 2001
From: ceci3 <ceci3@users.noreply.github.com>
Date: Fri, 12 Jun 2020 11:50:38 +0800
Subject: [PATCH] Fix warning (#345)

* fix

* refine warning
---
 paddleslim/common/rl_controller/__init__.py | 4 +---
 paddleslim/nas/rl_nas.py                    | 9 +++++++++
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/paddleslim/common/rl_controller/__init__.py b/paddleslim/common/rl_controller/__init__.py
index da815888..b29c5637 100644
--- a/paddleslim/common/rl_controller/__init__.py
+++ b/paddleslim/common/rl_controller/__init__.py
@@ -19,9 +19,7 @@ try:
     import parl
     from .ddpg import *
 except ImportError as e:
-    _logger.warn(
-        "If you want to use DDPG in RLNAS, please pip install parl first. Now states: {}".
-        format(e))
+    pass
 
 from .lstm import *
 from .utils import *
diff --git a/paddleslim/nas/rl_nas.py b/paddleslim/nas/rl_nas.py
index a7fa6591..1718b834 100644
--- a/paddleslim/nas/rl_nas.py
+++ b/paddleslim/nas/rl_nas.py
@@ -76,6 +76,15 @@ class RLNAS(object):
         self.save_controller = save_controller
         self.load_controller = load_controller
 
+        if key.upper() in ['DDPG']:
+            try:
+                import parl
+            except ImportError as e:
+                _logger.error(
+                    "If you want to use DDPG in RLNAS, please pip install parl first. Now states: {}".
+                    format(e))
+                os._exit(1)
+
         cls = RLCONTROLLER.get(key.upper())
 
         server_ip, server_port = server_addr
-- 
GitLab