Skip to content

GitLab

Unverified Commit 3ee96695 authored by Dhruv Nair's avatar Dhruv Nair Committed by GitHub
Browse files

Allow Mochi Transformer to be split across multiple GPUs (#10300) 

update
parent 9764f229
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
src/diffusers/models/transformers/transformer_mochi.py
View file @ 3ee96695
...@@ -335,6 +335,7 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOri ...@@ -335,6 +335,7 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOri
""" """
_supports_gradient_checkpointing = True _supports_gradient_checkpointing = True
_no_split_modules = ["MochiTransformerBlock"]
@register_to_config @register_to_config
def __init__( def __init__(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment