diff --git a/sdk/python/kubeflow/storage_init_container/hugging_face.py b/sdk/python/kubeflow/storage_init_container/hugging_face.py
index 62ace28a86..a238f83ceb 100644
--- a/sdk/python/kubeflow/storage_init_container/hugging_face.py
+++ b/sdk/python/kubeflow/storage_init_container/hugging_face.py
@@ -1,10 +1,8 @@
 from kubeflow.storage_init_container.abstract_model_provider import modelProvider
 from kubeflow.storage_init_container.abstract_dataset_provider import datasetProvider
 from dataclasses import dataclass, field
-from typing import Literal
 from urllib.parse import urlparse
 import json, os
-from typing import Dict, Any, Union
 from datasets import load_dataset
 from peft import LoraConfig
 import transformers
diff --git a/sdk/python/kubeflow/training/api/training_client.py b/sdk/python/kubeflow/training/api/training_client.py
index 2599ac6330..3d59e109fe 100644
--- a/sdk/python/kubeflow/training/api/training_client.py
+++ b/sdk/python/kubeflow/training/api/training_client.py
@@ -16,7 +16,7 @@
 import logging
 import time
 import json
-from typing import Optional, Callable, List, Dict, Any, Set, Literal, Union
+from typing import Optional, Callable, List, Dict, Any, Set, Union
 import queue
 from kubernetes import client, config, watch
 
@@ -102,7 +102,7 @@ def train(
         namespace: str = None,
         num_workers: int = 1,
         num_procs_per_worker: int = 1,
-        storage_config: Dict[Literal["size", "storage_class"], str] = {"size": "10Gi"},
+        storage_config: Dict[str, str] = {"size": "10Gi", "storage_class": None},
         model_provider_parameters: HuggingFaceModelParams = None,
         dataset_provider_parameters: Union[HfDatasetParams, S3DatasetParams] = None,
         train_parameters: HuggingFaceTrainParams = None,