[fbsync] Add missing _version to the MLPBlock (#6113)

Summary: * Add missing `_version` to the MLPBlock * fix linter Reviewed By: NicolasHug Differential Revision: D36931838 fbshipit-source-id: f1871bca9aa4b2fb9fb9eb1a1a2c60f46f5fe674
pytorch · Jun 6, 2022 · d90fe85 · d90fe85
1 parent ed27879
commit d90fe85
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torchvision/models/vision_transformer.py b/torchvision/models/vision_transformer.py
@@ -40,6 +40,8 @@ class ConvStemConfig(NamedTuple):
 class MLPBlock(MLP):
     """Transformer MLP block."""
 
+    _version = 2
+
     def __init__(self, in_dim: int, mlp_dim: int, dropout: float):
         super().__init__(in_dim, [mlp_dim, in_dim], activation_layer=nn.GELU, inplace=None, dropout=dropout)