Skip to content

GitLab

Unverified Commit f9e957f0 authored by Aryan's avatar Aryan Committed by GitHub
Browse files

Fix offload tests for CogVideoX and CogView3 (#10547) 

* update

* update
parent 4dec63c1
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
tests/models/transformers/test_models_transformer_cogvideox.py
View file @ f9e957f0
...@@ -33,6 +33,7 @@ class CogVideoXTransformerTests(ModelTesterMixin, unittest.TestCase): ...@@ -33,6 +33,7 @@ class CogVideoXTransformerTests(ModelTesterMixin, unittest.TestCase):
model_class = CogVideoXTransformer3DModel model_class = CogVideoXTransformer3DModel
main_input_name = "hidden_states" main_input_name = "hidden_states"
uses_custom_attn_processor = True uses_custom_attn_processor = True
model_split_percents = [0.7, 0.7, 0.8]
@property @property
def dummy_input(self): def dummy_input(self):
......
tests/models/transformers/test_models_transformer_cogview3plus.py
View file @ f9e957f0
...@@ -33,6 +33,7 @@ class CogView3PlusTransformerTests(ModelTesterMixin, unittest.TestCase): ...@@ -33,6 +33,7 @@ class CogView3PlusTransformerTests(ModelTesterMixin, unittest.TestCase):
model_class = CogView3PlusTransformer2DModel model_class = CogView3PlusTransformer2DModel
main_input_name = "hidden_states" main_input_name = "hidden_states"
uses_custom_attn_processor = True uses_custom_attn_processor = True
model_split_percents = [0.7, 0.6, 0.6]
@property @property
def dummy_input(self): def dummy_input(self):
......
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