Fix padding in dreambooth (#1030)

2022-11-03 00:37:05 +09:00 · 2022-11-03 00:37:05 +09:00 · 33c487455e
parent 5cd29d623a
commit 33c487455e
1 changed files with 6 additions and 1 deletions
--- 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,