diff --git a/learnware/market/heterogeneous/organizer/hetero_map/__init__.py b/learnware/market/heterogeneous/organizer/hetero_map/__init__.py
index 4ebb0a8..e124bcb 100644
--- a/learnware/market/heterogeneous/organizer/hetero_map/__init__.py
+++ b/learnware/market/heterogeneous/organizer/hetero_map/__init__.py
@@ -9,7 +9,7 @@ from torch import Tensor, nn
 from .....utils import allocate_cuda_idx, choose_device
 from .....specification import HeteroMapTableSpecification, RKMETableSpecification
 from .feature_extractor import CLSToken, FeatureProcessor, FeatureTokenizer
-from .trainer import TransTabCollatorForCL
+from .trainer import TransTabCollatorForCL, Trainer
 
 
 class HeteroMap(nn.Module):