diff --git a/comfy/sd.py b/comfy/sd.py
index edd0b51d8..b24498bd8 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -668,7 +668,8 @@ def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, m
 
     model_management.load_models_gpu(load_models, force_patch_weights=True)
     clip_vision_sd = clip_vision.get_sd() if clip_vision is not None else None
-    sd = model.model.state_dict_for_saving(clip_sd, vae.get_sd(), clip_vision_sd)
+    vae = vae.get_sd() if vae is not None else None
+    sd = model.model.state_dict_for_saving(clip_sd, vae, clip_vision_sd)
     for k in extra_keys:
         sd[k] = extra_keys[k]