diff --git a/traceml/traceml/exceptions.py b/traceml/traceml/exceptions.py
index 365318a00..5d6f4e900 100644
--- a/traceml/traceml/exceptions.py
+++ b/traceml/traceml/exceptions.py
@@ -1,5 +1,3 @@
 from polyaxon.exceptions import PolyaxonClientException
 
-
-class TracemlException(PolyaxonClientException):
-    pass
+TracemlException = PolyaxonClientException
diff --git a/traceml/traceml/integrations/tensorboard.py b/traceml/traceml/integrations/tensorboard.py
index f50c83856..f0beef6f7 100644
--- a/traceml/traceml/integrations/tensorboard.py
+++ b/traceml/traceml/integrations/tensorboard.py
@@ -2,7 +2,6 @@
 
 from typing import TYPE_CHECKING, Any, Optional
 
-from polyaxon.exceptions import PolyaxonClientException
 from traceml import tracking
 from traceml.exceptions import TracemlException
 from traceml.logger import logger
@@ -64,9 +63,7 @@ def process_summary(
                     log_histo=log_histo,
                     log_tensor=log_tensor,
                 )
-            except PolyaxonClientException(
-                "Polyaxon failed processing tensorboard summary."
-            ):
+            except TracemlException("Polyaxon failed processing tensorboard summary."):
                 pass
 
     @classmethod