Merge pull request modelscope#519 from XDUWQ/fix_custom

fix a bug of custom diffusion
liuyhwangyh · Sep 4, 2023 · 4345f3a · 4345f3a
2 parents 98f65f6 + 8d1c290
commit 4345f3a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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: