fix memory consistency decoder test (#5828)

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>

fix memory consistency decoder test (#5828)
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
2a84e8bb · Will Berman · GitHub · c896b841 · 2a84e8bb
Unverified Commit 2a84e8bb authored Nov 17, 2023 by Will Berman Committed by GitHub Nov 17, 2023
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

tests/models/test_models_vae.py tests/models/test_models_vae.py +1 -0

No files found.
--- a/tests/models/test_models_vae.py
+++ b/tests/models/test_models_vae.py
@@ -804,6 +804,7 @@ class ConsistencyDecoderVAEIntegrationTests(unittest.TestCase):
        gc.collect()
        torch.cuda.empty_cache()

+    @torch.no_grad()
    def test_encode_decode(self):
        vae = ConsistencyDecoderVAE.from_pretrained("openai/consistency-decoder")  # TODO - update
        vae.to(torch_device)