Skip to content

Commit 9727cda

Browse files
authored
[Tests] Mark the ncsnpp model tests as slow (#575)
* [Tests] Mark the ncsnpp model tests as slow * style
1 parent 0a2c42f commit 9727cda

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

tests/test_models_unet.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
import torch
2020

2121
from diffusers import UNet2DModel
22-
from diffusers.testing_utils import floats_tensor, torch_device
22+
from diffusers.testing_utils import floats_tensor, slow, torch_device
2323

2424
from .test_modeling_common import ModelTesterMixin
2525

@@ -231,6 +231,7 @@ def prepare_init_args_and_inputs_for_common(self):
231231
inputs_dict = self.dummy_input
232232
return init_dict, inputs_dict
233233

234+
@slow
234235
def test_from_pretrained_hub(self):
235236
model, loading_info = UNet2DModel.from_pretrained("google/ncsnpp-celebahq-256", output_loading_info=True)
236237
self.assertIsNotNone(model)
@@ -244,6 +245,7 @@ def test_from_pretrained_hub(self):
244245

245246
assert image is not None, "Make sure output is not None"
246247

248+
@slow
247249
def test_output_pretrained_ve_mid(self):
248250
model = UNet2DModel.from_pretrained("google/ncsnpp-celebahq-256")
249251
model.to(torch_device)

0 commit comments

Comments
 (0)