Skip to content

Commit f9e957f

Browse files
authored
Fix offload tests for CogVideoX and CogView3 (huggingface#10547)
* update * update
1 parent 4dec63c commit f9e957f

File tree

2 files changed

+2
-0
lines changed

2 files changed

+2
-0
lines changed

tests/models/transformers/test_models_transformer_cogvideox.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ class CogVideoXTransformerTests(ModelTesterMixin, unittest.TestCase):
3333
model_class = CogVideoXTransformer3DModel
3434
main_input_name = "hidden_states"
3535
uses_custom_attn_processor = True
36+
model_split_percents = [0.7, 0.7, 0.8]
3637

3738
@property
3839
def dummy_input(self):

tests/models/transformers/test_models_transformer_cogview3plus.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ class CogView3PlusTransformerTests(ModelTesterMixin, unittest.TestCase):
3333
model_class = CogView3PlusTransformer2DModel
3434
main_input_name = "hidden_states"
3535
uses_custom_attn_processor = True
36+
model_split_percents = [0.7, 0.6, 0.6]
3637

3738
@property
3839
def dummy_input(self):

0 commit comments

Comments
 (0)