Fix common tests for FP16 (huggingface#1588)

* Fix common tests for FP16 * revert
kno10 · Dec 7, 2022 · dc87f52 · dc87f52
1 parent d9b5b43
commit dc87f52
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/tests/pipelines/ddim/test_ddim.py b/tests/pipelines/ddim/test_ddim.py
@@ -29,6 +29,7 @@
 
 class DDIMPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
     pipeline_class = DDIMPipeline
+    test_cpu_offload = False
 
     def get_dummy_components(self):
         torch.manual_seed(0)

diff --git a/tests/test_pipelines_common.py b/tests/test_pipelines_common.py
@@ -190,7 +190,7 @@ def test_save_load_float16(self):
 
         with tempfile.TemporaryDirectory() as tmpdir:
             pipe.save_pretrained(tmpdir)
-            pipe_loaded = self.pipeline_class.from_pretrained(tmpdir)
+            pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, torch_dtype=torch.float16)
             pipe_loaded.to(torch_device)
             pipe_loaded.set_progress_bar_config(disable=None)