Fix issue with Krea2 patch size when training with shift

This commit is contained in:
Jaret Burkett committed 2026-09-16 13:46:39 -06:00
1 parent 410206e41c
commit e65c4d0fb6
1 file changed
+5
@@ -121,6 +121,11 @@ class SingleMMDiTConfig:
txtheads: int = 20
txtkvheads: int = 20
@property
def patch_size(self) -> int:
# trainer probes config.patch_size for the dynamic-shift token count
return self.patch
class SimpleModulation(torch.nn.Module):
def __init__(self, dim: int):