diff --git a/traceml/traceml/tracking/__init__.py b/traceml/traceml/tracking/__init__.py
index a66cdc434..25b8aa08e 100644
--- a/traceml/traceml/tracking/__init__.py
+++ b/traceml/traceml/tracking/__init__.py
@@ -1,7 +1,8 @@
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Sequence, Union
 
-from polyaxon.client import *
+from polyaxon._client.run import RunClient
+from polyaxon._sdk.schemas import V1Run
 from traceml.artifacts import V1ArtifactKind, V1RunArtifact
 from traceml.tracking.run import Run