Fix padding in dreambooth (open-mmlab#1030)

aravind-h-v · Nov 2, 2022 · 33c4874 · 33c4874
1 parent 5cd29d6
commit 33c4874
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py
@@ -494,7 +494,12 @@ def collate_fn(examples):
         pixel_values = torch.stack(pixel_values)
         pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()
 
-        input_ids = tokenizer.pad({"input_ids": input_ids}, padding=True, return_tensors="pt").input_ids
+        input_ids = tokenizer.pad(
+            {"input_ids": input_ids},
+            padding="max_length",
+            max_length=tokenizer.model_max_length,
+            return_tensors="pt",
+        ).input_ids
 
         batch = {
             "input_ids": input_ids,