Skip to content

GitLab

Unverified Commit 2a84e8bb authored by Will Berman's avatar Will Berman Committed by GitHub
Browse files

fix memory consistency decoder test (#5828) 


Co-authored-by: default avatarDhruv Nair <dhruv.nair@gmail.com>
parent c896b841
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
tests/models/test_models_vae.py
View file @ 2a84e8bb
...@@ -804,6 +804,7 @@ class ConsistencyDecoderVAEIntegrationTests(unittest.TestCase): ...@@ -804,6 +804,7 @@ class ConsistencyDecoderVAEIntegrationTests(unittest.TestCase):
gc.collect() gc.collect()
torch.cuda.empty_cache() torch.cuda.empty_cache()
@torch.no_grad()
def test_encode_decode(self): def test_encode_decode(self):
vae = ConsistencyDecoderVAE.from_pretrained("openai/consistency-decoder") # TODO - update vae = ConsistencyDecoderVAE.from_pretrained("openai/consistency-decoder") # TODO - update
vae.to(torch_device) vae.to(torch_device)
......
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