We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Add Pytorch-Image-Models to backbones
## 🚀 Feature
https://github.com/rwightman/pytorch-image-models#models

### Motivation

The torch vision models are very limiting these models seem to perform best on kaggle challenges would be great to integrate them 

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `flash/vision/backbones.py`
Content:
```
1 # Copyright The PyTorch Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 from typing import Tuple
15 
16 import torchvision
17 from pytorch_lightning.utilities import _BOLTS_AVAILABLE, rank_zero_warn
18 from pytorch_lightning.utilities.exceptions import MisconfigurationException
19 from torch import nn as nn
20 from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
21 
22 if _BOLTS_AVAILABLE:
23     from pl_bolts.models.self_supervised import SimCLR, SwAV
24 
25 ROOT_S3_BUCKET = "https://pl-bolts-weights.s3.us-east-2.amazonaws.com"
26 
27 MOBILENET_MODELS = ["mobilenet_v2"]
28 VGG_MODELS = ["vgg11", "vgg13", "vgg16", "vgg19"]
29 RESNET_MODELS = ["resnet18", "resnet34", "resnet50", "resnet101", "resnet152", "resnext50_32x4d", "resnext101_32x8d"]
30 DENSENET_MODELS = ["densenet121", "densenet169", "densenet161"]
31 TORCHVISION_MODELS = MOBILENET_MODELS + VGG_MODELS + RESNET_MODELS + DENSENET_MODELS
32 
33 BOLTS_MODELS = ["simclr-imagenet", "swav-imagenet"]
34 
35 
36 def backbone_and_num_features(
37     model_name: str,
38     fpn: bool = False,
39     pretrained: bool = True,
40     trainable_backbone_layers: int = 3,
41     **kwargs
42 ) -> Tuple[nn.Module, int]:
43     """
44     Args:
45         model_name: backbone supported by `torchvision` and `bolts`
46         fpn: If True, creates a Feature Pyramind Network on top of Resnet based CNNs.
47         pretrained: if true, returns a model with backbone pre-trained on Imagenet
48         trainable_backbone_layers: number of trainable resnet layers starting from final block.
49 
50     >>> backbone_and_num_features('mobilenet_v2')  # doctest: +ELLIPSIS
51     (Sequential(...), 1280)
52     >>> backbone_and_num_features('resnet50', fpn=True)  # doctest: +ELLIPSIS
53     (BackboneWithFPN(...), 256)
54     >>> backbone_and_num_features('swav-imagenet')  # doctest: +ELLIPSIS
55     (Sequential(...), 2048)
56     """
57     if fpn:
58         if model_name in RESNET_MODELS:
59             backbone = resnet_fpn_backbone(
60                 model_name, pretrained=pretrained, trainable_layers=trainable_backbone_layers, **kwargs
61             )
62             fpn_out_channels = 256
63             return backbone, fpn_out_channels
64         else:
65             rank_zero_warn(f"{model_name} backbone is not supported with `fpn=True`, `fpn` won't be added.")
66 
67     if model_name in BOLTS_MODELS:
68         return bolts_backbone_and_num_features(model_name)
69 
70     if model_name in TORCHVISION_MODELS:
71         return torchvision_backbone_and_num_features(model_name, pretrained)
72 
73     raise ValueError(f"{model_name} is not supported yet.")
74 
75 
76 def bolts_backbone_and_num_features(model_name: str) -> Tuple[nn.Module, int]:
77     """
78     >>> bolts_backbone_and_num_features('simclr-imagenet')  # doctest: +ELLIPSIS
79     (Sequential(...), 2048)
80     >>> bolts_backbone_and_num_features('swav-imagenet')  # doctest: +ELLIPSIS
81     (Sequential(...), 2048)
82     """
83 
84     # TODO: maybe we should plain pytorch weights so we don't need to rely on bolts to load these
85     # also mabye just use torchhub for the ssl lib
86     def load_simclr_imagenet(path_or_url: str = f"{ROOT_S3_BUCKET}/simclr/bolts_simclr_imagenet/simclr_imagenet.ckpt"):
87         simclr = SimCLR.load_from_checkpoint(path_or_url, strict=False)
88         # remove the last two layers & turn it into a Sequential model
89         backbone = nn.Sequential(*list(simclr.encoder.children())[:-2])
90         return backbone, 2048
91 
92     def load_swav_imagenet(path_or_url: str = f"{ROOT_S3_BUCKET}/swav/swav_imagenet/swav_imagenet.pth.tar"):
93         swav = SwAV.load_from_checkpoint(path_or_url, strict=True)
94         # remove the last two layers & turn it into a Sequential model
95         backbone = nn.Sequential(*list(swav.model.children())[:-2])
96         return backbone, 2048
97 
98     models = {
99         'simclr-imagenet': load_simclr_imagenet,
100         'swav-imagenet': load_swav_imagenet,
101     }
102     if not _BOLTS_AVAILABLE:
103         raise MisconfigurationException("Bolts isn't installed. Please, use ``pip install lightning-bolts``.")
104     if model_name in models:
105         return models[model_name]()
106 
107     raise ValueError(f"{model_name} is not supported yet.")
108 
109 
110 def torchvision_backbone_and_num_features(model_name: str, pretrained: bool = True) -> Tuple[nn.Module, int]:
111     """
112     >>> torchvision_backbone_and_num_features('mobilenet_v2')  # doctest: +ELLIPSIS
113     (Sequential(...), 1280)
114     >>> torchvision_backbone_and_num_features('resnet18')  # doctest: +ELLIPSIS
115     (Sequential(...), 512)
116     >>> torchvision_backbone_and_num_features('densenet121')  # doctest: +ELLIPSIS
117     (Sequential(...), 1024)
118     """
119     model = getattr(torchvision.models, model_name, None)
120     if model is None:
121         raise MisconfigurationException(f"{model_name} is not supported by torchvision")
122 
123     if model_name in MOBILENET_MODELS + VGG_MODELS:
124         model = model(pretrained=pretrained)
125         backbone = model.features
126         num_features = 512 if model_name in VGG_MODELS else model.classifier[-1].in_features
127         return backbone, num_features
128 
129     elif model_name in RESNET_MODELS:
130         model = model(pretrained=pretrained)
131         # remove the last two layers & turn it into a Sequential model
132         backbone = nn.Sequential(*list(model.children())[:-2])
133         num_features = model.fc.in_features
134         return backbone, num_features
135 
136     elif model_name in DENSENET_MODELS:
137         model = model(pretrained=pretrained)
138         backbone = nn.Sequential(*model.features, nn.ReLU(inplace=True))
139         num_features = model.classifier.in_features
140         return backbone, num_features
141 
142     raise ValueError(f"{model_name} is not supported yet.")
143
```

Path: `flash/utils/imports.py`
Content:
```
1 from pytorch_lightning.utilities.imports import _module_available
2 
3 _TABNET_AVAILABLE = _module_available("pytorch_tabnet")
4 _KORNIA_AVAILABLE = _module_available("kornia")
5 _COCO_AVAILABLE = _module_available("pycocotools")
6
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/flash/utils/imports.py b/flash/utils/imports.py
--- a/flash/utils/imports.py
+++ b/flash/utils/imports.py
@@ -3,3 +3,4 @@
 _TABNET_AVAILABLE = _module_available("pytorch_tabnet")
 _KORNIA_AVAILABLE = _module_available("kornia")
 _COCO_AVAILABLE = _module_available("pycocotools")
+_TIMM_AVAILABLE = _module_available("timm")
diff --git a/flash/vision/backbones.py b/flash/vision/backbones.py
--- a/flash/vision/backbones.py
+++ b/flash/vision/backbones.py
@@ -19,6 +19,11 @@
 from torch import nn as nn
 from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
 
+from flash.utils.imports import _TIMM_AVAILABLE
+
+if _TIMM_AVAILABLE:
+    import timm
+
 if _BOLTS_AVAILABLE:
     from pl_bolts.models.self_supervised import SimCLR, SwAV
 
@@ -70,6 +75,9 @@
     if model_name in TORCHVISION_MODELS:
         return torchvision_backbone_and_num_features(model_name, pretrained)
 
+    if _TIMM_AVAILABLE and model_name in timm.list_models():
+        return timm_backbone_and_num_features(model_name, pretrained)
+
     raise ValueError(f"{model_name} is not supported yet.")
 
 
@@ -140,3 +148,15 @@
         return backbone, num_features
 
     raise ValueError(f"{model_name} is not supported yet.")
+
+
+def timm_backbone_and_num_features(model_name: str, pretrained: bool = True) -> Tuple[nn.Module, int]:
+
+    if model_name in timm.list_models():
+        backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0, global_pool='')
+        num_features = backbone.num_features
+        return backbone, num_features
+
+    raise ValueError(
+        f"{model_name} is not supported in timm yet. https://rwightman.github.io/pytorch-image-models/models/"
+    )

{"golden_diff": "diff --git a/flash/utils/imports.py b/flash/utils/imports.py\n--- a/flash/utils/imports.py\n+++ b/flash/utils/imports.py\n@@ -3,3 +3,4 @@\n _TABNET_AVAILABLE = _module_available(\"pytorch_tabnet\")\n _KORNIA_AVAILABLE = _module_available(\"kornia\")\n _COCO_AVAILABLE = _module_available(\"pycocotools\")\n+_TIMM_AVAILABLE = _module_available(\"timm\")\ndiff --git a/flash/vision/backbones.py b/flash/vision/backbones.py\n--- a/flash/vision/backbones.py\n+++ b/flash/vision/backbones.py\n@@ -19,6 +19,11 @@\n from torch import nn as nn\n from torchvision.models.detection.backbone_utils import resnet_fpn_backbone\n \n+from flash.utils.imports import _TIMM_AVAILABLE\n+\n+if _TIMM_AVAILABLE:\n+    import timm\n+\n if _BOLTS_AVAILABLE:\n     from pl_bolts.models.self_supervised import SimCLR, SwAV\n \n@@ -70,6 +75,9 @@\n     if model_name in TORCHVISION_MODELS:\n         return torchvision_backbone_and_num_features(model_name, pretrained)\n \n+    if _TIMM_AVAILABLE and model_name in timm.list_models():\n+        return timm_backbone_and_num_features(model_name, pretrained)\n+\n     raise ValueError(f\"{model_name} is not supported yet.\")\n \n \n@@ -140,3 +148,15 @@\n         return backbone, num_features\n \n     raise ValueError(f\"{model_name} is not supported yet.\")\n+\n+\n+def timm_backbone_and_num_features(model_name: str, pretrained: bool = True) -> Tuple[nn.Module, int]:\n+\n+    if model_name in timm.list_models():\n+        backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0, global_pool='')\n+        num_features = backbone.num_features\n+        return backbone, num_features\n+\n+    raise ValueError(\n+        f\"{model_name} is not supported in timm yet. https://rwightman.github.io/pytorch-image-models/models/\"\n+    )\n", "issue": "Add Pytorch-Image-Models to backbones\n## \ud83d\ude80 Feature\r\nhttps://github.com/rwightman/pytorch-image-models#models\r\n\r\n### Motivation\r\n\r\nThe torch vision models are very limiting these models seem to perform best on kaggle challenges would be great to integrate them \n", "before_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Tuple\n\nimport torchvision\nfrom pytorch_lightning.utilities import _BOLTS_AVAILABLE, rank_zero_warn\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\nfrom torch import nn as nn\nfrom torchvision.models.detection.backbone_utils import resnet_fpn_backbone\n\nif _BOLTS_AVAILABLE:\n    from pl_bolts.models.self_supervised import SimCLR, SwAV\n\nROOT_S3_BUCKET = \"https://pl-bolts-weights.s3.us-east-2.amazonaws.com\"\n\nMOBILENET_MODELS = [\"mobilenet_v2\"]\nVGG_MODELS = [\"vgg11\", \"vgg13\", \"vgg16\", \"vgg19\"]\nRESNET_MODELS = [\"resnet18\", \"resnet34\", \"resnet50\", \"resnet101\", \"resnet152\", \"resnext50_32x4d\", \"resnext101_32x8d\"]\nDENSENET_MODELS = [\"densenet121\", \"densenet169\", \"densenet161\"]\nTORCHVISION_MODELS = MOBILENET_MODELS + VGG_MODELS + RESNET_MODELS + DENSENET_MODELS\n\nBOLTS_MODELS = [\"simclr-imagenet\", \"swav-imagenet\"]\n\n\ndef backbone_and_num_features(\n    model_name: str,\n    fpn: bool = False,\n    pretrained: bool = True,\n    trainable_backbone_layers: int = 3,\n    **kwargs\n) -> Tuple[nn.Module, int]:\n    \"\"\"\n    Args:\n        model_name: backbone supported by `torchvision` and `bolts`\n        fpn: If True, creates a Feature Pyramind Network on top of Resnet based CNNs.\n        pretrained: if true, returns a model with backbone pre-trained on Imagenet\n        trainable_backbone_layers: number of trainable resnet layers starting from final block.\n\n    >>> backbone_and_num_features('mobilenet_v2')  # doctest: +ELLIPSIS\n    (Sequential(...), 1280)\n    >>> backbone_and_num_features('resnet50', fpn=True)  # doctest: +ELLIPSIS\n    (BackboneWithFPN(...), 256)\n    >>> backbone_and_num_features('swav-imagenet')  # doctest: +ELLIPSIS\n    (Sequential(...), 2048)\n    \"\"\"\n    if fpn:\n        if model_name in RESNET_MODELS:\n            backbone = resnet_fpn_backbone(\n                model_name, pretrained=pretrained, trainable_layers=trainable_backbone_layers, **kwargs\n            )\n            fpn_out_channels = 256\n            return backbone, fpn_out_channels\n        else:\n            rank_zero_warn(f\"{model_name} backbone is not supported with `fpn=True`, `fpn` won't be added.\")\n\n    if model_name in BOLTS_MODELS:\n        return bolts_backbone_and_num_features(model_name)\n\n    if model_name in TORCHVISION_MODELS:\n        return torchvision_backbone_and_num_features(model_name, pretrained)\n\n    raise ValueError(f\"{model_name} is not supported yet.\")\n\n\ndef bolts_backbone_and_num_features(model_name: str) -> Tuple[nn.Module, int]:\n    \"\"\"\n    >>> bolts_backbone_and_num_features('simclr-imagenet')  # doctest: +ELLIPSIS\n    (Sequential(...), 2048)\n    >>> bolts_backbone_and_num_features('swav-imagenet')  # doctest: +ELLIPSIS\n    (Sequential(...), 2048)\n    \"\"\"\n\n    # TODO: maybe we should plain pytorch weights so we don't need to rely on bolts to load these\n    # also mabye just use torchhub for the ssl lib\n    def load_simclr_imagenet(path_or_url: str = f\"{ROOT_S3_BUCKET}/simclr/bolts_simclr_imagenet/simclr_imagenet.ckpt\"):\n        simclr = SimCLR.load_from_checkpoint(path_or_url, strict=False)\n        # remove the last two layers & turn it into a Sequential model\n        backbone = nn.Sequential(*list(simclr.encoder.children())[:-2])\n        return backbone, 2048\n\n    def load_swav_imagenet(path_or_url: str = f\"{ROOT_S3_BUCKET}/swav/swav_imagenet/swav_imagenet.pth.tar\"):\n        swav = SwAV.load_from_checkpoint(path_or_url, strict=True)\n        # remove the last two layers & turn it into a Sequential model\n        backbone = nn.Sequential(*list(swav.model.children())[:-2])\n        return backbone, 2048\n\n    models = {\n        'simclr-imagenet': load_simclr_imagenet,\n        'swav-imagenet': load_swav_imagenet,\n    }\n    if not _BOLTS_AVAILABLE:\n        raise MisconfigurationException(\"Bolts isn't installed. Please, use ``pip install lightning-bolts``.\")\n    if model_name in models:\n        return models[model_name]()\n\n    raise ValueError(f\"{model_name} is not supported yet.\")\n\n\ndef torchvision_backbone_and_num_features(model_name: str, pretrained: bool = True) -> Tuple[nn.Module, int]:\n    \"\"\"\n    >>> torchvision_backbone_and_num_features('mobilenet_v2')  # doctest: +ELLIPSIS\n    (Sequential(...), 1280)\n    >>> torchvision_backbone_and_num_features('resnet18')  # doctest: +ELLIPSIS\n    (Sequential(...), 512)\n    >>> torchvision_backbone_and_num_features('densenet121')  # doctest: +ELLIPSIS\n    (Sequential(...), 1024)\n    \"\"\"\n    model = getattr(torchvision.models, model_name, None)\n    if model is None:\n        raise MisconfigurationException(f\"{model_name} is not supported by torchvision\")\n\n    if model_name in MOBILENET_MODELS + VGG_MODELS:\n        model = model(pretrained=pretrained)\n        backbone = model.features\n        num_features = 512 if model_name in VGG_MODELS else model.classifier[-1].in_features\n        return backbone, num_features\n\n    elif model_name in RESNET_MODELS:\n        model = model(pretrained=pretrained)\n        # remove the last two layers & turn it into a Sequential model\n        backbone = nn.Sequential(*list(model.children())[:-2])\n        num_features = model.fc.in_features\n        return backbone, num_features\n\n    elif model_name in DENSENET_MODELS:\n        model = model(pretrained=pretrained)\n        backbone = nn.Sequential(*model.features, nn.ReLU(inplace=True))\n        num_features = model.classifier.in_features\n        return backbone, num_features\n\n    raise ValueError(f\"{model_name} is not supported yet.\")\n", "path": "flash/vision/backbones.py"}, {"content": "from pytorch_lightning.utilities.imports import _module_available\n\n_TABNET_AVAILABLE = _module_available(\"pytorch_tabnet\")\n_KORNIA_AVAILABLE = _module_available(\"kornia\")\n_COCO_AVAILABLE = _module_available(\"pycocotools\")\n", "path": "flash/utils/imports.py"}], "after_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Tuple\n\nimport torchvision\nfrom pytorch_lightning.utilities import _BOLTS_AVAILABLE, rank_zero_warn\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\nfrom torch import nn as nn\nfrom torchvision.models.detection.backbone_utils import resnet_fpn_backbone\n\nfrom flash.utils.imports import _TIMM_AVAILABLE\n\nif _TIMM_AVAILABLE:\n    import timm\n\nif _BOLTS_AVAILABLE:\n    from pl_bolts.models.self_supervised import SimCLR, SwAV\n\nROOT_S3_BUCKET = \"https://pl-bolts-weights.s3.us-east-2.amazonaws.com\"\n\nMOBILENET_MODELS = [\"mobilenet_v2\"]\nVGG_MODELS = [\"vgg11\", \"vgg13\", \"vgg16\", \"vgg19\"]\nRESNET_MODELS = [\"resnet18\", \"resnet34\", \"resnet50\", \"resnet101\", \"resnet152\", \"resnext50_32x4d\", \"resnext101_32x8d\"]\nDENSENET_MODELS = [\"densenet121\", \"densenet169\", \"densenet161\"]\nTORCHVISION_MODELS = MOBILENET_MODELS + VGG_MODELS + RESNET_MODELS + DENSENET_MODELS\n\nBOLTS_MODELS = [\"simclr-imagenet\", \"swav-imagenet\"]\n\n\ndef backbone_and_num_features(\n    model_name: str,\n    fpn: bool = False,\n    pretrained: bool = True,\n    trainable_backbone_layers: int = 3,\n    **kwargs\n) -> Tuple[nn.Module, int]:\n    \"\"\"\n    Args:\n        model_name: backbone supported by `torchvision` and `bolts`\n        fpn: If True, creates a Feature Pyramind Network on top of Resnet based CNNs.\n        pretrained: if true, returns a model with backbone pre-trained on Imagenet\n        trainable_backbone_layers: number of trainable resnet layers starting from final block.\n\n    >>> backbone_and_num_features('mobilenet_v2')  # doctest: +ELLIPSIS\n    (Sequential(...), 1280)\n    >>> backbone_and_num_features('resnet50', fpn=True)  # doctest: +ELLIPSIS\n    (BackboneWithFPN(...), 256)\n    >>> backbone_and_num_features('swav-imagenet')  # doctest: +ELLIPSIS\n    (Sequential(...), 2048)\n    \"\"\"\n    if fpn:\n        if model_name in RESNET_MODELS:\n            backbone = resnet_fpn_backbone(\n                model_name, pretrained=pretrained, trainable_layers=trainable_backbone_layers, **kwargs\n            )\n            fpn_out_channels = 256\n            return backbone, fpn_out_channels\n        else:\n            rank_zero_warn(f\"{model_name} backbone is not supported with `fpn=True`, `fpn` won't be added.\")\n\n    if model_name in BOLTS_MODELS:\n        return bolts_backbone_and_num_features(model_name)\n\n    if model_name in TORCHVISION_MODELS:\n        return torchvision_backbone_and_num_features(model_name, pretrained)\n\n    if _TIMM_AVAILABLE and model_name in timm.list_models():\n        return timm_backbone_and_num_features(model_name, pretrained)\n\n    raise ValueError(f\"{model_name} is not supported yet.\")\n\n\ndef bolts_backbone_and_num_features(model_name: str) -> Tuple[nn.Module, int]:\n    \"\"\"\n    >>> bolts_backbone_and_num_features('simclr-imagenet')  # doctest: +ELLIPSIS\n    (Sequential(...), 2048)\n    >>> bolts_backbone_and_num_features('swav-imagenet')  # doctest: +ELLIPSIS\n    (Sequential(...), 2048)\n    \"\"\"\n\n    # TODO: maybe we should plain pytorch weights so we don't need to rely on bolts to load these\n    # also mabye just use torchhub for the ssl lib\n    def load_simclr_imagenet(path_or_url: str = f\"{ROOT_S3_BUCKET}/simclr/bolts_simclr_imagenet/simclr_imagenet.ckpt\"):\n        simclr = SimCLR.load_from_checkpoint(path_or_url, strict=False)\n        # remove the last two layers & turn it into a Sequential model\n        backbone = nn.Sequential(*list(simclr.encoder.children())[:-2])\n        return backbone, 2048\n\n    def load_swav_imagenet(path_or_url: str = f\"{ROOT_S3_BUCKET}/swav/swav_imagenet/swav_imagenet.pth.tar\"):\n        swav = SwAV.load_from_checkpoint(path_or_url, strict=True)\n        # remove the last two layers & turn it into a Sequential model\n        backbone = nn.Sequential(*list(swav.model.children())[:-2])\n        return backbone, 2048\n\n    models = {\n        'simclr-imagenet': load_simclr_imagenet,\n        'swav-imagenet': load_swav_imagenet,\n    }\n    if not _BOLTS_AVAILABLE:\n        raise MisconfigurationException(\"Bolts isn't installed. Please, use ``pip install lightning-bolts``.\")\n    if model_name in models:\n        return models[model_name]()\n\n    raise ValueError(f\"{model_name} is not supported yet.\")\n\n\ndef torchvision_backbone_and_num_features(model_name: str, pretrained: bool = True) -> Tuple[nn.Module, int]:\n    \"\"\"\n    >>> torchvision_backbone_and_num_features('mobilenet_v2')  # doctest: +ELLIPSIS\n    (Sequential(...), 1280)\n    >>> torchvision_backbone_and_num_features('resnet18')  # doctest: +ELLIPSIS\n    (Sequential(...), 512)\n    >>> torchvision_backbone_and_num_features('densenet121')  # doctest: +ELLIPSIS\n    (Sequential(...), 1024)\n    \"\"\"\n    model = getattr(torchvision.models, model_name, None)\n    if model is None:\n        raise MisconfigurationException(f\"{model_name} is not supported by torchvision\")\n\n    if model_name in MOBILENET_MODELS + VGG_MODELS:\n        model = model(pretrained=pretrained)\n        backbone = model.features\n        num_features = 512 if model_name in VGG_MODELS else model.classifier[-1].in_features\n        return backbone, num_features\n\n    elif model_name in RESNET_MODELS:\n        model = model(pretrained=pretrained)\n        # remove the last two layers & turn it into a Sequential model\n        backbone = nn.Sequential(*list(model.children())[:-2])\n        num_features = model.fc.in_features\n        return backbone, num_features\n\n    elif model_name in DENSENET_MODELS:\n        model = model(pretrained=pretrained)\n        backbone = nn.Sequential(*model.features, nn.ReLU(inplace=True))\n        num_features = model.classifier.in_features\n        return backbone, num_features\n\n    raise ValueError(f\"{model_name} is not supported yet.\")\n\n\ndef timm_backbone_and_num_features(model_name: str, pretrained: bool = True) -> Tuple[nn.Module, int]:\n\n    if model_name in timm.list_models():\n        backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0, global_pool='')\n        num_features = backbone.num_features\n        return backbone, num_features\n\n    raise ValueError(\n        f\"{model_name} is not supported in timm yet. https://rwightman.github.io/pytorch-image-models/models/\"\n    )\n", "path": "flash/vision/backbones.py"}, {"content": "from pytorch_lightning.utilities.imports import _module_available\n\n_TABNET_AVAILABLE = _module_available(\"pytorch_tabnet\")\n_KORNIA_AVAILABLE = _module_available(\"kornia\")\n_COCO_AVAILABLE = _module_available(\"pycocotools\")\n_TIMM_AVAILABLE = _module_available(\"timm\")\n", "path": "flash/utils/imports.py"}]}