-skip

hlky · hlky · commit 5f889332882f · 2024-12-18T11:53:10.000Z
diff --git a/tests/models/unets/test_models_unet_2d.py b/tests/models/unets/test_models_unet_2d.py
@@ -361,10 +361,6 @@ def test_forward_with_norm_groups(self):
     def test_gradient_checkpointing_is_applied(self):
         expected_set = {
             "UNetMidBlock2D",
-            "SkipDownBlock2D",
-            "AttnSkipDownBlock2D",
-            "SkipUpBlock2D",
-            "AttnSkipUpBlock2D",
         }
 
         block_out_channels = (32, 64, 64, 64)
