Skip to content

Commit 99200be

Browse files
committed
address comment
1 parent e38675b commit 99200be

1 file changed

Lines changed: 2 additions & 2 deletions

File tree

src/maxdiffusion/train_flux.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,10 +42,10 @@ def main(argv: Sequence[str]) -> None:
4242

4343
@contextmanager
4444
def transformer_engine_context():
45-
""" If TransformerEngine is available, this context manager will provide the library with MaxText-specific details needed for correcct operation. """
45+
""" If TransformerEngine is available, this context manager will provide the library with MaxDiffusion-specific details needed for correcct operation. """
4646
try:
4747
from transformer_engine.jax.sharding import global_shard_guard, MeshResource
48-
# Inform TransformerEngine of MaxText's physical mesh resources.
48+
# Inform TransformerEngine of MaxDiffusion's physical mesh resources.
4949
mesh_resource = MeshResource(
5050
dp_resource = "data",
5151
tp_resource = "tensor",

0 commit comments

Comments
 (0)