diff --git a/model_zoo/official/nlp/tinybert/run_general_distill.py b/model_zoo/official/nlp/tinybert/run_general_distill.py
index 36b88487df42a804cb961f09fb5b6df8daf4e406..1b64ed4dfbb7940a85e2a6a651ea620b0cfe61a1 100644
--- a/model_zoo/official/nlp/tinybert/run_general_distill.py
+++ b/model_zoo/official/nlp/tinybert/run_general_distill.py
@@ -106,7 +106,7 @@ def run_general_distill():
 
     if args_opt.dataset_type == "tfrecord":
         dataset_type = DataType.TFRECORD
-    elif arg_opt.dataset_type == "mindrecord":
+    elif args_opt.dataset_type == "mindrecord":
         dataset_type = DataType.MINDRECORD
     else:
         raise Exception("dataset format is not supported yet")
diff --git a/model_zoo/official/nlp/tinybert/run_task_distill.py b/model_zoo/official/nlp/tinybert/run_task_distill.py
index e0e78c4ed039a191fbe02d1682fc9311a25d4ab2..0b93a2b3bd32e3aac8973fed44bdd3d76085f582 100644
--- a/model_zoo/official/nlp/tinybert/run_task_distill.py
+++ b/model_zoo/official/nlp/tinybert/run_task_distill.py
@@ -125,9 +125,9 @@ def run_predistill():
     rank = 0
     device_num = 1
 
-    if arg_opt.dataset_type == "tfrecord":
+    if args_opt.dataset_type == "tfrecord":
         dataset_type = DataType.TFRECORD
-    elif arg_opt.dataset_type == "mindrecord":
+    elif args_opt.dataset_type == "mindrecord":
         dataset_type = DataType.MINDRECORD
     else:
         raise Exception("dataset format is not supported yet")