Fix kontext finetune issue when batch size >1 (#11921)

mymusise · web-flow · commit cde02b061b6f · 2025-07-18T19:38:58.000-04:00
set drop_last to True

Signed-off-by: mymusise &lt;mymusise1@gmail.com&gt;
diff --git a/examples/dreambooth/train_dreambooth_lora_flux_kontext.py b/examples/dreambooth/train_dreambooth_lora_flux_kontext.py
@@ -1614,7 +1614,7 @@ def load_model_hook(models, input_dir):
     )
     if args.cond_image_column is not None:
         logger.info("I2I fine-tuning enabled.")
-    batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=False)
+    batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=True)
     train_dataloader = torch.utils.data.DataLoader(
         train_dataset,
         batch_sampler=batch_sampler,

Original file line number	Diff line number	Diff line change
`@@ -1614,7 +1614,7 @@ def load_model_hook(models, input_dir):`
`1614`	`1614`	`)`
`1615`	`1615`	`if args.cond_image_column is not None:`
`1616`	`1616`	`logger.info("I2I fine-tuning enabled.")`
`1617`		`- batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=False)`
	`1617`	`+ batch_sampler = BucketBatchSampler(train_dataset, batch_size=args.train_batch_size, drop_last=True)`
`1618`	`1618`	`train_dataloader = torch.utils.data.DataLoader(`
`1619`	`1619`	`train_dataset,`
`1620`	`1620`	`batch_sampler=batch_sampler,`