diff --git a/modules/text/text_generation/unified_transformer-12L-cn-luge/module.py b/modules/text/text_generation/unified_transformer-12L-cn-luge/module.py
index 52ef5532db84d696960d4ac28ef1cba4bbfbc75c..115b1e0e9ea102332bfb35009fa8d474276ba5a6 100644
--- a/modules/text/text_generation/unified_transformer-12L-cn-luge/module.py
+++ b/modules/text/text_generation/unified_transformer-12L-cn-luge/module.py
@@ -54,7 +54,7 @@ class UnifiedTransformer(nn.Layer):
         Generate input batches.
         """
         padding = False if batch_size == 1 else True
-        pad_func = Pad(pad_val=self.tokenizer.pad_token_id, pad_right=False)
+        pad_func = Pad(pad_val=self.tokenizer.pad_token_id, pad_right=False, dtype=np.int64)
 
         def pad_mask(batch_attention_mask):
             batch_size = len(batch_attention_mask)
@@ -75,9 +75,9 @@ class UnifiedTransformer(nn.Layer):
                 position_ids = pad_func([example['position_ids'] for example in batch_examples])
                 attention_mask = pad_mask([example['attention_mask'] for example in batch_examples])
             else:
-                input_ids = np.asarray([example['input_ids'] for example in batch_examples])
-                token_type_ids = np.asarray([example['token_type_ids'] for example in batch_examples])
-                position_ids = np.asarray([example['position_ids'] for example in batch_examples])
+                input_ids = np.asarray([example['input_ids'] for example in batch_examples], dtype=np.int64)
+                token_type_ids = np.asarray([example['token_type_ids'] for example in batch_examples], dtype=np.int64)
+                position_ids = np.asarray([example['position_ids'] for example in batch_examples], dtype=np.int64)
                 attention_mask = np.asarray([example['attention_mask'] for example in batch_examples])
                 attention_mask = np.expand_dims(attention_mask, 0)
 
diff --git a/modules/text/text_generation/unified_transformer-12L-cn/module.py b/modules/text/text_generation/unified_transformer-12L-cn/module.py
index ee09a55d0c2853a7abfcf6a19bc727c1de5c1ad2..363d15d709aff6c314844634a97080bcb28f57c5 100644
--- a/modules/text/text_generation/unified_transformer-12L-cn/module.py
+++ b/modules/text/text_generation/unified_transformer-12L-cn/module.py
@@ -54,7 +54,7 @@ class UnifiedTransformer(nn.Layer):
         Generate input batches.
         """
         padding = False if batch_size == 1 else True
-        pad_func = Pad(pad_val=self.tokenizer.pad_token_id, pad_right=False)
+        pad_func = Pad(pad_val=self.tokenizer.pad_token_id, pad_right=False, dtype=np.int64)
 
         def pad_mask(batch_attention_mask):
             batch_size = len(batch_attention_mask)
@@ -75,9 +75,9 @@ class UnifiedTransformer(nn.Layer):
                 position_ids = pad_func([example['position_ids'] for example in batch_examples])
                 attention_mask = pad_mask([example['attention_mask'] for example in batch_examples])
             else:
-                input_ids = np.asarray([example['input_ids'] for example in batch_examples])
-                token_type_ids = np.asarray([example['token_type_ids'] for example in batch_examples])
-                position_ids = np.asarray([example['position_ids'] for example in batch_examples])
+                input_ids = np.asarray([example['input_ids'] for example in batch_examples], dtype=np.int64)
+                token_type_ids = np.asarray([example['token_type_ids'] for example in batch_examples], dtype=np.int64)
+                position_ids = np.asarray([example['position_ids'] for example in batch_examples], dtype=np.int64)
                 attention_mask = np.asarray([example['attention_mask'] for example in batch_examples])
                 attention_mask = np.expand_dims(attention_mask, 0)