Repository: pytorch/vision

Issue:
Fix flaky `test_classification_model[cuda-vit_h_14]` test
### 🐛 Describe the bug

After the addition of `vit_h_14` (#5210) we started receiving memory issues in some of our tests. Unfortunately the don't fail consistently but rather occasionally. Examples [1](https://app.circleci.com/pipelines/github/pytorch/vision/14208/workflows/be78ecf7-48b2-4c9a-be99-505930f86e35/jobs/1143437) and [2](https://app.circleci.com/pipelines/github/pytorch/vision/14208/workflows/be78ecf7-48b2-4c9a-be99-505930f86e35/jobs/1143479). 

Sometimes the memory issues are limited to that specific test, while other (see example 2) it takes down with it all other tests. We need to address this and restore our CI to green.

### Versions

Latest main branch - 4bf6c6e4b4cf2dea9c7f9952aa30d5820ab8ae33


Discussion:

Resolving patch:
diff --git a/hubconf.py b/hubconf.py
--- a/hubconf.py
+++ b/hubconf.py
@@ -63,5 +63,4 @@
     vit_b_32,
     vit_l_16,
     vit_l_32,
-    vit_h_14,
 )
diff --git a/torchvision/models/vision_transformer.py b/torchvision/models/vision_transformer.py
--- a/torchvision/models/vision_transformer.py
+++ b/torchvision/models/vision_transformer.py
@@ -15,7 +15,6 @@
     "vit_b_32",
     "vit_l_16",
     "vit_l_32",
-    "vit_h_14",
 ]
 
 model_urls = {
@@ -357,26 +356,6 @@ def vit_l_32(pretrained: bool = False, progress: bool = True, **kwargs: Any) ->
     )
 
 
-def vit_h_14(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VisionTransformer:
-    """
-    Constructs a vit_h_14 architecture from
-    `"An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" <https://arxiv.org/abs/2010.11929>`_.
-
-    NOTE: Pretrained weights are not available for this model.
-    """
-    return _vision_transformer(
-        arch="vit_h_14",
-        patch_size=14,
-        num_layers=32,
-        num_heads=16,
-        hidden_dim=1280,
-        mlp_dim=5120,
-        pretrained=pretrained,
-        progress=progress,
-        **kwargs,
-    )
-
-
 def interpolate_embeddings(
     image_size: int,
     patch_size: int,
diff --git a/torchvision/prototype/models/vision_transformer.py b/torchvision/prototype/models/vision_transformer.py
--- a/torchvision/prototype/models/vision_transformer.py
+++ b/torchvision/prototype/models/vision_transformer.py
@@ -19,12 +19,10 @@
     "ViT_B_32_Weights",
     "ViT_L_16_Weights",
     "ViT_L_32_Weights",
-    "ViT_H_14_Weights",
     "vit_b_16",
     "vit_b_32",
     "vit_l_16",
     "vit_l_32",
-    "vit_h_14",
 ]
 
 
@@ -105,11 +103,6 @@ class ViT_L_32_Weights(WeightsEnum):
     default = ImageNet1K_V1
 
 
-class ViT_H_14_Weights(WeightsEnum):
-    # Weights are not available yet.
-    pass
-
-
 def _vision_transformer(
     patch_size: int,
     num_layers: int,
@@ -203,19 +196,3 @@ def vit_l_32(*, weights: Optional[ViT_L_32_Weights] = None, progress: bool = Tru
         progress=progress,
         **kwargs,
     )
-
-
-@handle_legacy_interface(weights=("pretrained", None))
-def vit_h_14(*, weights: Optional[ViT_H_14_Weights] = None, progress: bool = True, **kwargs: Any) -> VisionTransformer:
-    weights = ViT_H_14_Weights.verify(weights)
-
-    return _vision_transformer(
-        patch_size=14,
-        num_layers=32,
-        num_heads=16,
-        hidden_dim=1280,
-        mlp_dim=5120,
-        weights=weights,
-        progress=progress,
-        **kwargs,
-    )


Test patch:
diff --git a/test/expect/ModelTester.test_vit_h_14_expect.pkl b/test/expect/ModelTester.test_vit_h_14_expect.pkl
deleted file mode 100644
Binary files a/test/expect/ModelTester.test_vit_h_14_expect.pkl and /dev/null differ