fix kvae gradient checkpointing tests (#13865)

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>

Author: sayakpaul

tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py

@@ -97,6 +97,9 @@ 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, "
@@ -105,13 +108,13 @@ def test_training_with_ema(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."""