Skip to content

GitLab

Unverified Commit 1d1f648c authored by Leon Lin's avatar Leon Lin Committed by GitHub
Browse files

fix dreambooth attention mask (#3541)

parent 67cf0445
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
examples/dreambooth/train_dreambooth.py
View file @ 1d1f648c
...@@ -681,6 +681,7 @@ def collate_fn(examples, with_prior_preservation=False): ...@@ -681,6 +681,7 @@ def collate_fn(examples, with_prior_preservation=False):
} }
if has_attention_mask: if has_attention_mask:
attention_mask = torch.cat(attention_mask, dim=0)
batch["attention_mask"] = attention_mask batch["attention_mask"] = attention_mask
return batch return batch
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment