diff --git a/modelscope/trainers/multi_modal/custom_diffusion/custom_diffusion_trainer.py b/modelscope/trainers/multi_modal/custom_diffusion/custom_diffusion_trainer.py
index 1183c167b..a18b546e6 100644
--- a/modelscope/trainers/multi_modal/custom_diffusion/custom_diffusion_trainer.py
+++ b/modelscope/trainers/multi_modal/custom_diffusion/custom_diffusion_trainer.py
@@ -40,7 +40,8 @@ class CustomCheckpointProcessor(CheckpointProcessor):
     def __init__(self,
                  modifier_token,
                  modifier_token_id,
-                 torch_type=torch.float32):
+                 torch_type=torch.float32,
+                 safe_serialization=False):
         """Checkpoint processor for custom diffusion.
 
         Args: