fix nas's bug (PaddlePaddle#1019)

lizexu123 · Mar 21, 2022 · d66afeb · d66afeb
1 parent 1a9376a
commit d66afeb
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/paddleslim/common/controller_client.py b/paddleslim/common/controller_client.py
@@ -17,7 +17,7 @@
 import logging
 import socket
 from .log_helper import get_logger
-
+import sys
 __all__ = ['ControllerClient']
 
 _logger = get_logger(__name__, level=logging.INFO)
@@ -110,7 +110,10 @@ def next_tokens(self):
 
         socket_client.send("next_tokens".encode())
         tokens = socket_client.recv(1024).decode()
-        tokens = [int(token) for token in tokens.strip("\n").split(",")]
+        try:
+            tokens = [int(token) for token in tokens.strip("\n").split(",")]
+        except ValueError:
+            sys.exit()
         return tokens
 
     def request_current_info(self):