Fix seed checkpoint creation assertion (#1112)

Fixes #1107

Author: rualark

torchtitan/experiments/flux/train.py

@@ -226,8 +226,8 @@ def eval_step(self, prompt: str = "A photo of a cat"):
     try:
         trainer = FluxTrainer(config)
         if config.checkpoint.create_seed_checkpoint:
-            assert int(
-                os.environ["WORLD_SIZE"]
+            assert (
+                int(os.environ["WORLD_SIZE"]) == 1
             ), "Must create seed checkpoint using a single device, to disable sharding."
             assert (
                 config.checkpoint.enable_checkpoint

torchtitan/train.py

@@ -455,8 +455,8 @@ def close(self) -> None:
         trainer = Trainer(config)
 
         if config.checkpoint.create_seed_checkpoint:
-            assert int(
-                os.environ["WORLD_SIZE"]
+            assert (
+                int(os.environ["WORLD_SIZE"]) == 1
             ), "Must create seed checkpoint using a single device, to disable sharding."
             assert (
                 config.checkpoint.enable_checkpoint