Revert "fix kvae gradient checkpointing tests"

This reverts commit 52ecdb0.

Author: sayakpaul

tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py

@@ -97,9 +97,6 @@ def test_training(self):
     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, "
@@ -108,13 +105,13 @@ def test_mixed_precision_training(self):
     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."""