Enable multi-process DataLoader for dreambooth (#950)

This commit is contained in:
Simon Kirsten 2022-10-26 16:24:48 +01:00 committed by GitHub
parent bd06dd023f
commit 8332c1a6d9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions

View File

@ -502,7 +502,7 @@ def main():
return batch
train_dataloader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.train_batch_size, shuffle=True, collate_fn=collate_fn
train_dataset, batch_size=args.train_batch_size, shuffle=True, collate_fn=collate_fn, num_workers=1
)
# Scheduler and math around the number of training steps.