diff --git a/autogen/oai/client.py b/autogen/oai/client.py
index 63a0923afe61..38af18eb0afc 100644
--- a/autogen/oai/client.py
+++ b/autogen/oai/client.py
@@ -274,7 +274,7 @@ def create(self, params: Dict[str, Any]) -> ChatCompletion:
             model = chunk.model.replace("gpt-35", "gpt-3.5")  # hack for Azure API
             try:
                 prompt_tokens = count_token(params["messages"], model)
-            except Exception as e:
+            except NotImplementedError as e:
                 # Catch token calculation error if streaming with customized models.
                 logger.warning(str(e))
                 prompt_tokens = 0