From 52ecdb07eb11ad19bdaeb47f28ee3f5d01492969 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Thu, 4 Jun 2026 14:12:51 +0530 Subject: [PATCH] fix kvae gradient checkpointing tests --- .../test_models_autoencoder_kl_kvae_video.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py index d810145c4d..dc74842db8 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py @@ -97,6 +97,9 @@ class TestAutoencoderKLKVAEVideoTraining(AutoencoderKLKVAEVideoTesterConfig, Tra def test_training_with_ema(self): _run_nondeterministic(super().test_training_with_ema) + def test_mixed_precision_training(self): + _run_nondeterministic(super().test_mixed_precision_training) + @pytest.mark.skip( "Gradient checkpointing recomputes the forward pass, but the model uses a stateful cache_dict " "that is mutated during the first forward. On recomputation the cache is already populated, " @@ -105,13 +108,13 @@ class TestAutoencoderKLKVAEVideoTraining(AutoencoderKLKVAEVideoTesterConfig, Tra def test_gradient_checkpointing_equivalence(self): super().test_gradient_checkpointing_equivalence() - def test_layerwise_casting_training(self): - _run_nondeterministic(super().test_layerwise_casting_training) - class TestAutoencoderKLKVAEVideoMemory(AutoencoderKLKVAEVideoTesterConfig, MemoryTesterMixin): """Memory optimization tests for AutoencoderKLKVAEVideo.""" + def test_layerwise_casting_training(self): + _run_nondeterministic(super().test_layerwise_casting_training) + class TestAutoencoderKLKVAEVideoSlicingTiling(AutoencoderKLKVAEVideoTesterConfig, NewAutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKLKVAEVideo."""