Skip to content

GitLab

Unverified Commit d36103a0 authored by Patrick von Platen's avatar Patrick von Platen Committed by GitHub
Browse files

[Tests] Speed up test (#2919) 

speed up test
parent b3c437e0
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 6 deletions
tests/models/test_models_unet_3d_condition.py
View file @ d36103a0
...@@ -88,19 +88,17 @@ class UNet3DConditionModelTests(ModelTesterMixin, unittest.TestCase): ...@@ -88,19 +88,17 @@ class UNet3DConditionModelTests(ModelTesterMixin, unittest.TestCase):
def prepare_init_args_and_inputs_for_common(self): def prepare_init_args_and_inputs_for_common(self):
init_dict = { init_dict = {
"block_out_channels": (32, 64, 64, 64), "block_out_channels": (32, 64),
"down_block_types": ( "down_block_types": (
"CrossAttnDownBlock3D",
"CrossAttnDownBlock3D",
"CrossAttnDownBlock3D", "CrossAttnDownBlock3D",
"DownBlock3D", "DownBlock3D",
), ),
"up_block_types": ("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D"), "up_block_types": ("UpBlock3D", "CrossAttnUpBlock3D"),
"cross_attention_dim": 32, "cross_attention_dim": 32,
"attention_head_dim": 4, "attention_head_dim": 8,
"out_channels": 4, "out_channels": 4,
"in_channels": 4, "in_channels": 4,
"layers_per_block": 2, "layers_per_block": 1,
"sample_size": 32, "sample_size": 32,
} }
inputs_dict = self.dummy_input inputs_dict = self.dummy_input
......
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