diff --git a/convlab/dst/trippy/modeling_dst.py b/convlab/dst/trippy/modeling_dst.py
index 2828d17ed1e97ebb60b74ab999c28efb1e7bfa88..3bd875b60b5059ae00888756a926080295af8a9b 100644
--- a/convlab/dst/trippy/modeling_dst.py
+++ b/convlab/dst/trippy/modeling_dst.py
@@ -62,7 +62,7 @@ def TransformerForDST(parent_name):
     class TransformerForDST(PARENT_CLASSES[parent_name]):
         def __init__(self, config):
             assert config.model_type in PARENT_CLASSES
-            assert self.__class__.__bases__[0] in MODEL_CLASSES
+            # assert self.__class__.__bases__[0] in MODEL_CLASSES
             super(TransformerForDST, self).__init__(config)
             self.model_type = config.model_type
             self.slot_list = config.dst_slot_list
@@ -82,7 +82,7 @@ def TransformerForDST(parent_name):
                 self.refer_index = -1
 
             # Make sure this module has the same name as in the pretrained checkpoint you want to load!
-            self.add_module(self.model_type, MODEL_CLASSES[self.__class__.__bases__[0]](config))
+            self.add_module(self.model_type, MODEL_CLASSES[PARENT_CLASSES[self.model_type]](config))
             if self.model_type == "electra":
                 self.pooler = ElectraPooler(config)
             
diff --git a/convlab/dst/trippy/tracker.py b/convlab/dst/trippy/tracker.py
index b0470266b2ce2c7d5cfba29d0270972b0c7cfa78..8ceaedddc5392c7a4ba9cdb7fe2f9ac9b39a72ba 100644
--- a/convlab/dst/trippy/tracker.py
+++ b/convlab/dst/trippy/tracker.py
@@ -30,10 +30,15 @@ from convlab.dst.trippy.modeling_dst import (TransformerForDST)
 from convlab.dst.trippy.dataset_interfacer import (create_dataset_interfacer)
 from convlab.util import relative_import_module_from_unified_datasets
 
+
+class BertForDST(TransformerForDST('bert')): pass
+class RobertaForDST(TransformerForDST('roberta')): pass
+class ElectraForDST(TransformerForDST('electra')): pass
+
 MODEL_CLASSES = {
-    'bert': (BertConfig, TransformerForDST('bert'), BertTokenizer),
-    'roberta': (RobertaConfig, TransformerForDST('roberta'), RobertaTokenizer),
-    'electra': (ElectraConfig, TransformerForDST('electra'), ElectraTokenizer),
+    'bert': (BertConfig, BertForDST, BertTokenizer),
+    'roberta': (RobertaConfig, RobertaForDST, RobertaTokenizer),
+    'electra': (ElectraConfig, ElectraForDST, ElectraTokenizer),
 }