Fix flux image_generation_test. Change input dim to match encoder (#1726)

yosoyjay · web-flow · commit 22d2d44de566 · 2025-09-22T23:50:40.000-07:00
Quick fix to address issue a failing test for `experiments/flux` that I encountered and opened and described in #1725. Fix is to match `context_in_dim` to the corresponding text embedding size of the encoder model `t5-v1_1-base`.
diff --git a/torchtitan/experiments/flux/tests/test_generate_image.py b/torchtitan/experiments/flux/tests/test_generate_image.py
@@ -98,7 +98,7 @@ def test_generate_image(self):
         t1 = time.perf_counter()
 
         model = self._get_test_model(
-            context_in_dim=4096, device=torch_device, dtype=torch.bfloat16
+            context_in_dim=768, device=torch_device, dtype=torch.bfloat16
         )
         model.eval()
 

Original file line number	Diff line number	Diff line change
`@@ -98,7 +98,7 @@ def test_generate_image(self):`
`98`	`98`	`t1 = time.perf_counter()`
`99`	`99`
`100`	`100`	`model = self._get_test_model(`
`101`		`- context_in_dim=4096, device=torch_device, dtype=torch.bfloat16`
	`101`	`+ context_in_dim=768, device=torch_device, dtype=torch.bfloat16`
`102`	`102`	`)`
`103`	`103`	`model.eval()`
`104`	`104`