Fix padding in dreambooth (#1030)

33c48745 · Yuta Hayashibe · GitHub · 5cd29d62 · 33c48745
Unverified Commit 33c48745 authored Nov 03, 2022 by Yuta Hayashibe Committed by GitHub Nov 02, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 1 deletion

examples/dreambooth/train_dreambooth.py examples/dreambooth/train_dreambooth.py +6 -1

No files found.
--- a/examples/dreambooth/train_dreambooth.py
+++ b/examples/dreambooth/train_dreambooth.py
@@ -494,7 +494,12 @@ def main(args):
        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,