Fix compatibility with VISSL master

Lightning-Universe · Apr 13, 2022 · 8db29e8 · 8db29e8
1 parent e51563b
commit 8db29e8
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/flash/image/embedding/backbones/vissl_backbones.py b/flash/image/embedding/backbones/vissl_backbones.py
@@ -60,6 +60,8 @@ def vision_transformer(
                     "QKV_BIAS": qkv_bias,
                     "QK_SCALE": qk_scale,
                     "CLASSIFIER": classifier,
+                    "CHECKPOINT_MLP": False,
+                    "CHECKPOINT_BLOCK": False,
                 }
             ),
         }