Change the file name

train_dit.py

@@ -165,9 +165,7 @@ def train(args):
                 "step": step_count,
                 "optimizer": optimizer.state_dict(),
             },
-            os.path.join(
-                train_config["task_name"], train_config["vae_autoencoder_ckpt_name"]
-            ),
+            os.path.join(train_config["task_name"], train_config["dit_ckpt_name"]),
         )
         if (epoch + 1) % 5 == 0:  # Save every 5 epochs
             artifact = wandb.Artifact(
@@ -176,9 +174,7 @@ def train(args):
                 description=f"DIT model checkpoint at epoch {epoch + 1}",
             )
             artifact.add_file(
-                os.path.join(
-                    train_config["task_name"], train_config["vae_autoencoder_ckpt_name"]
-                ),
+                os.path.join(train_config["task_name"], train_config["dit_ckpt_name"]),
             )
             wandb.log_artifact(artifact)
 
@@ -186,9 +182,7 @@ def train(args):
         "dit_model_final", type="model", description="Final DIT model checkpoint"
     )
     final_artifact.add_file(
-        os.path.join(
-            train_config["task_name"], train_config["vae_autoencoder_ckpt_name"]
-        )
+        os.path.join(train_config["task_name"], train_config["dit_ckpt_name"])
     )
     wandb.log_artifact(final_artifact)