diff --git a/PaddleNLP/examples/machine_reading_comprehension/DuReader-robust/run_du.py b/PaddleNLP/examples/machine_reading_comprehension/DuReader-robust/run_du.py
index feaed689391bd1de76f9c7acb477379f61bdd30c..6b2614db3bf5fd5ed5abd3a939906dd7aed92d65 100644
--- a/PaddleNLP/examples/machine_reading_comprehension/DuReader-robust/run_du.py
+++ b/PaddleNLP/examples/machine_reading_comprehension/DuReader-robust/run_du.py
@@ -181,7 +181,7 @@ def do_train(args):
         args.learning_rate,
         lambda current_step, warmup_proportion=args.warmup_proportion,
         num_training_steps=args.max_steps if args.max_steps > 0 else
-        (len(train_ds.examples)//args.batch_size*args.num_train_epochs): float(
+        (len(train_data_loader)*args.num_train_epochs): float(
             current_step) / float(max(1, warmup_proportion*num_training_steps))
         if current_step < warmup_proportion*num_training_steps else max(
             0.0,
diff --git a/PaddleNLP/paddlenlp/datasets/squad.py b/PaddleNLP/paddlenlp/datasets/squad.py
index f5355aa44a737f8e553e62b0e6e546abc105a6f0..467d5293902741d7ecb273729f9539cfd685e959 100644
--- a/PaddleNLP/paddlenlp/datasets/squad.py
+++ b/PaddleNLP/paddlenlp/datasets/squad.py
@@ -243,7 +243,14 @@ class SQuAD(Dataset):
                 segment_ids.append(1)
 
                 input_ids = tokenizer.convert_tokens_to_ids(tokens)
-
+                input_ids = input_ids + [
+                    tokenizer.vocab[tokenizer.pad_token]
+                    for _ in range(self.max_seq_length - len(input_ids))
+                ]
+                segment_ids = segment_ids + [
+                    tokenizer.vocab[tokenizer.pad_token]
+                    for _ in range(self.max_seq_length - len(segment_ids))
+                ]
                 input_mask = [1] * len(input_ids)
 
                 start_position = None