[Flax] Fix sample batch size DreamBooth (open-mmlab#1129)

fix sample batch size
aravind-h-v · Nov 4, 2022 · c62b3a2 · c62b3a2
1 parent bde4880
commit c62b3a2
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/examples/dreambooth/train_dreambooth_flax.py b/examples/dreambooth/train_dreambooth_flax.py
@@ -361,7 +361,8 @@ def main():
             logger.info(f"Number of class images to sample: {num_new_images}.")
 
             sample_dataset = PromptDataset(args.class_prompt, num_new_images)
-            sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=args.sample_batch_size)
+            total_sample_batch_size = args.sample_batch_size * jax.local_device_count()
+            sample_dataloader = torch.utils.data.DataLoader(sample_dataset, batch_size=total_sample_batch_size)
 
             for example in tqdm(
                 sample_dataloader, desc="Generating class images", disable=not jax.process_index() == 0