From 2a84e8bb5a30905566b0ae718de9beeecc97e420 Mon Sep 17 00:00:00 2001
From: Will Berman <wlbberman@gmail.com>
Date: Fri, 17 Nov 2023 09:31:01 -0800
Subject: [PATCH] fix memory consistency decoder test (#5828)

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
---
 tests/models/test_models_vae.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/models/test_models_vae.py b/tests/models/test_models_vae.py
index 3b698624ff..5283639269 100644
--- 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)