ltx_transformer_step_test.py modified

Author: prishajain1

src/maxdiffusion/tests/ltx_transformer_step_test.py

@@ -108,7 +108,7 @@ def test_one_step_transformer(self):
 
     with open(config_path, "r") as f:
       model_config = json.load(f)
-    relative_ckpt_path = model_config["ckpt_path"]
+    relative_ckpt_path = model_config.get("ckpt_path")
     ignored_keys = [
         "_class_name",
         "_diffusers_version",
@@ -129,9 +129,11 @@ def test_one_step_transformer(self):
         transformer.init_weights, in_channels, key, model_config["caption_channels"], eval_only=True
     )
 
-    absolute_ckpt_path = os.path.abspath(relative_ckpt_path)
-
-    checkpoint_manager = ocp.CheckpointManager(absolute_ckpt_path)
+    if relative_ckpt_path:
+      absolute_ckpt_path = os.path.abspath(relative_ckpt_path)
+      checkpoint_manager = ocp.CheckpointManager(absolute_ckpt_path)
+    else:
+      checkpoint_manager = None
     transformer_state, transformer_state_shardings = setup_initial_state(
         model=transformer,
         tx=None,