diff --git a/requirements.txt b/requirements.txt
index d4d06d1..312cab3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,4 @@
 # NOTE: once we add more dependencies, consider update dependabot to check for updates
 
-lightning-sdk >=0.1.35
+lightning-sdk >=0.1.40
 lightning-utilities
diff --git a/src/litmodels/io/cloud.py b/src/litmodels/io/cloud.py
index be81b27..ad31b00 100644
--- a/src/litmodels/io/cloud.py
+++ b/src/litmodels/io/cloud.py
@@ -2,13 +2,15 @@
 # Licensed under the Apache License, Version 2.0 (the "License");
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
-from typing import Optional, Tuple, Union
+from typing import TYPE_CHECKING, Optional, Tuple, Union
 
-from lightning_sdk.api.teamspace_api import UploadedModelInfo
 from lightning_sdk.lightning_cloud.env import LIGHTNING_CLOUD_URL
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.utils import resolve as sdk_resolvers
 
+if TYPE_CHECKING:
+    from lightning_sdk.models import UploadedModelInfo
+
 # if module_available("lightning"):
 #     from lightning import LightningModule
 # elif module_available("pytorch_lightning"):
@@ -81,7 +83,7 @@ def upload_model_files(
     progress_bar: bool = True,
     cluster_id: Optional[str] = None,
     verbose: Union[bool, int] = 1,
-) -> UploadedModelInfo:
+) -> "UploadedModelInfo":
     """Upload a local checkpoint file to the model store.
 
     Args:
diff --git a/src/litmodels/io/gateway.py b/src/litmodels/io/gateway.py
index d2d89d7..cd437c0 100644
--- a/src/litmodels/io/gateway.py
+++ b/src/litmodels/io/gateway.py
@@ -1,9 +1,8 @@
 import os
 import tempfile
 from pathlib import Path
-from typing import Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
-from lightning_sdk.api.teamspace_api import UploadedModelInfo
 from lightning_utilities import module_available
 
 from litmodels.io.cloud import download_model_files, upload_model_files
@@ -14,6 +13,9 @@
 else:
     torch = None
 
+if TYPE_CHECKING:
+    from lightning_sdk.models import UploadedModelInfo
+
 
 def upload_model(
     name: str,
@@ -22,7 +24,7 @@ def upload_model(
     cluster_id: Optional[str] = None,
     staging_dir: Optional[str] = None,
     verbose: Union[bool, int] = 1,
-) -> UploadedModelInfo:
+) -> "UploadedModelInfo":
     """Upload a checkpoint to the model store.
 
     Args: