update

a-r-r-o-w · a-r-r-o-w · commit 178c22dc51df · 2024-12-23T11:20:37.000+01:00
diff --git a/tests/models/autoencoders/test_models_autoencoder_ltx_video.py b/tests/models/autoencoders/test_models_autoencoder_ltx_video.py
@@ -82,11 +82,11 @@ def prepare_init_args_and_inputs_for_common(self):
 
     def test_gradient_checkpointing_is_applied(self):
         expected_set = {
-            "LTXEncoder3d",
-            "LTXDecoder3d",
-            "LTXDownBlock3D",
-            "LTXMidBlock3d",
-            "LTXUpBlock3d",
+            "LTXVideoEncoder3d",
+            "LTXVideoDecoder3d",
+            "LTXVideoDownBlock3D",
+            "LTXVideoMidBlock3d",
+            "LTXVideoUpBlock3d",
         }
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
@@ -152,11 +152,11 @@ def prepare_init_args_and_inputs_for_common(self):
 
     def test_gradient_checkpointing_is_applied(self):
         expected_set = {
-            "LTXEncoder3d",
-            "LTXDecoder3d",
-            "LTXDownBlock3D",
-            "LTXMidBlock3d",
-            "LTXUpBlock3d",
+            "LTXVideoEncoder3d",
+            "LTXVideoDecoder3d",
+            "LTXVideoDownBlock3D",
+            "LTXVideoMidBlock3d",
+            "LTXVideoUpBlock3d",
         }
         super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
