Fix: bf16 casting after checkpoint load for Stage 2

train.py

@@ -385,7 +385,8 @@ def train_stage2(args, config):
     if args.checkpoint:
         ckpt_path = os.path.join(args.checkpoint, "model.pt")
         if os.path.exists(ckpt_path):
-            model.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
+            state = torch.load(ckpt_path, map_location="cpu")
+            model.load_state_dict(state, strict=False)
             print(f"Loaded Stage 1 checkpoint from {ckpt_path}")
         else:
             print("No Stage 1 checkpoint found, training from scratch")
@@ -393,6 +394,10 @@ def train_stage2(args, config):
     else:
         model = initialize_mla_from_pretrained(model, config.base_model, config)
     
+    # Cast to bf16 AFTER loading checkpoint (ckpt weights may be fp32)
+    model = model.to(torch.bfloat16)
+    print("Model cast to bfloat16")
+    
     # Dataset
     dataset = ImageTextDataset(
         tokenizer, vae,