@@ -22,7 +22,7 @@ try:
     from deci_lab_client.client import DeciPlatformClient
     from deci_common.data_interfaces.files_data_interface import FilesDataInterface
     from deci_lab_client.models import AutoNACFileName
-    from deci_lab_client import ApiException, ModelMetadata, OptimizationRequestForm
+    from deci_lab_client import ApiException
 except (ImportError, NameError):
     client_enabled = False
 
@@ -131,14 +131,14 @@ class DeciClient:
                 f"you can override this by passing models.get(... download_required_code=False) and importing the files yourself"
             )
 
-    def upload_model(self, model: nn.Module, model_meta_data: ModelMetadata, optimization_request_form: OptimizationRequestForm):
+    def upload_model(self, model: nn.Module, model_meta_data, optimization_request_form):
         """
         This function will upload the trained model to the Deci Lab
 
         Args:
-            model: The resulting model from the training process
-            model_meta_data: metadata to accompany the model
-            optimization_request_form: the optimization parameters
+            model:                     The resulting model from the training process
+            model_meta_data:           Metadata to accompany the model
+            optimization_request_form: The optimization parameters
         """
         self.lab_client.login(token=os.getenv("DECI_PLATFORM_TOKEN"))
         self.lab_client.add_model(