We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
ImageEmbedder default behavior is not a flattened output
## 🐛 Bug

I discovered this issue while testing PR #655. If you run the [Image Embedding README example code](https://github.com/PyTorchLightning/lightning-flash#example-1-image-embedding), it returns a 3D tensor. 
My understanding from the use of embeddings in general, and how they are used in [Fifty One](https://voxel51.com/docs/fiftyone/tutorials/image_embeddings.html) is they expect the embeddings to be 1D (for each embedding). 

The reason it returns a 3D tensor is because it depends on the backbone used. The default there is `resnet101`, which returns a `2048x7x7` shape tensor. Others like inception return a flat 1D tensor, i.e. length-X.

### To Reproduce

Steps to reproduce the behavior:

Run the [README example](https://github.com/PyTorchLightning/lightning-flash#example-1-image-embedding), but remove the `embedding_dim` parameter. See below for example.

Note: as-is, this will error on `print(embeddings.shape)`, regardless of configuration, since that is a list. But the question here is around the logic for the ImageEmbedder. 


#### Code sample
```python
from flash.core.data.utils import download_data
from flash.image import ImageEmbedder

# 1. Download the data
download_data("https://pl-flash-data.s3.amazonaws.com/hymenoptera_data.zip", "data/")

# 2. Create an ImageEmbedder with resnet50 trained on imagenet.
embedder = ImageEmbedder(backbone="resnet50")

# 3. Generate an embedding from an image path.
embeddings = embedder.predict("data/hymenoptera_data/predict/153783656_85f9c3ac70.jpg")

# 4. Print embeddings shape
print(embeddings.shape)
```

### Expected behavior

Expect to see a 100352x1 shape tensor as the output, instead of 2048x7x7. 

### Environment

 - PyTorch Version (e.g., 1.0): 1.9
 - OS (e.g., Linux): Linux
 - How you installed PyTorch (`conda`, `pip`, source): pip
 - Build command you used (if compiling from source): N/A
 - Python version: 3.8.6
 - CUDA/cuDNN version: N/A
 - GPU models and configuration: N/A
 - Any other relevant information: N/A

### Additional context

I believe the question is around what the logic should be here:
https://github.com/PyTorchLightning/lightning-flash/blob/075de3a46d74d9fc0e769401063fede1f12d0518/flash/image/embedding/model.py#L85-L92

If `embedding_dim` is None, then the head is `nn.Identity()`. **If we desire a flat 1D embedding, then the question is: should `nn.Identity()` change to `nn.Flatten()`?**

It could be argued that the user should be left to flatten after on their own, but per the contributing guidelines, I thought this would align with "[Force User Decisions To Best Practices](https://github.com/PyTorchLightning/lightning-flash/blob/ddd942d3dfe3884a97a855446410166c3c9f16d9/.github/CONTRIBUTING.md#force-user-decisions-to-best-practices)"

Let me know your thoughts. If that makes sense, then I can update the code, run some tests, and update docs in a PR. 


--- 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/image/embedding/model.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 Any, Callable, Mapping, Optional, Sequence, Tuple, Type, Union
15 
16 import torch
17 from pytorch_lightning.utilities import rank_zero_warn
18 from torch import nn
19 from torch.nn import functional as F
20 from torchmetrics import Accuracy, Metric
21 
22 from flash.core.data.data_source import DefaultDataKeys
23 from flash.core.model import Task
24 from flash.core.registry import FlashRegistry
25 from flash.core.utilities.imports import _IMAGE_AVAILABLE
26 from flash.core.utilities.isinstance import _isinstance
27 from flash.image.classification.data import ImageClassificationPreprocess
28 
29 if _IMAGE_AVAILABLE:
30     from flash.image.classification.backbones import IMAGE_CLASSIFIER_BACKBONES
31 else:
32     IMAGE_CLASSIFIER_BACKBONES = FlashRegistry("backbones")
33 
34 
35 class ImageEmbedder(Task):
36     """The ``ImageEmbedder`` is a :class:`~flash.Task` for obtaining feature vectors (embeddings) from images. For
37     more details, see :ref:`image_embedder`.
38 
39     Args:
40         embedding_dim: Dimension of the embedded vector. ``None`` uses the default from the backbone.
41         backbone: A model to use to extract image features, defaults to ``"swav-imagenet"``.
42         pretrained: Use a pretrained backbone, defaults to ``True``.
43         loss_fn: Loss function for training and finetuning, defaults to :func:`torch.nn.functional.cross_entropy`
44         optimizer: Optimizer to use for training and finetuning, defaults to :class:`torch.optim.SGD`.
45         metrics: Metrics to compute for training and evaluation. Can either be an metric from the `torchmetrics`
46             package, a custom metric inherenting from `torchmetrics.Metric`, a callable function or a list/dict
47             containing a combination of the aforementioned. In all cases, each metric needs to have the signature
48             `metric(preds,target)` and return a single scalar tensor. Defaults to :class:`torchmetrics.Accuracy`.
49         learning_rate: Learning rate to use for training, defaults to ``1e-3``.
50         pooling_fn: Function used to pool image to generate embeddings, defaults to :func:`torch.max`.
51     """
52 
53     backbones: FlashRegistry = IMAGE_CLASSIFIER_BACKBONES
54 
55     required_extras: str = "image"
56 
57     def __init__(
58         self,
59         embedding_dim: Optional[int] = None,
60         backbone: str = "resnet101",
61         pretrained: bool = True,
62         loss_fn: Callable = F.cross_entropy,
63         optimizer: Type[torch.optim.Optimizer] = torch.optim.SGD,
64         metrics: Union[Metric, Callable, Mapping, Sequence, None] = (Accuracy()),
65         learning_rate: float = 1e-3,
66         pooling_fn: Callable = torch.max,
67     ):
68         super().__init__(
69             model=None,
70             loss_fn=loss_fn,
71             optimizer=optimizer,
72             metrics=metrics,
73             learning_rate=learning_rate,
74             preprocess=ImageClassificationPreprocess(),
75         )
76 
77         self.save_hyperparameters()
78         self.backbone_name = backbone
79         self.embedding_dim = embedding_dim
80         assert pooling_fn in [torch.mean, torch.max]
81         self.pooling_fn = pooling_fn
82 
83         self.backbone, num_features = self.backbones.get(backbone)(pretrained=pretrained)
84 
85         if embedding_dim is None:
86             self.head = nn.Identity()
87         else:
88             self.head = nn.Sequential(
89                 nn.Flatten(),
90                 nn.Linear(num_features, embedding_dim),
91             )
92             rank_zero_warn("Adding linear layer on top of backbone. Remember to finetune first before using!")
93 
94     def apply_pool(self, x):
95         x = self.pooling_fn(x, dim=-1)
96         if _isinstance(x, Tuple[torch.Tensor, torch.Tensor]):
97             x = x[0]
98         x = self.pooling_fn(x, dim=-1)
99         if _isinstance(x, Tuple[torch.Tensor, torch.Tensor]):
100             x = x[0]
101         return x
102 
103     def forward(self, x) -> torch.Tensor:
104         x = self.backbone(x)
105 
106         # bolts ssl models return lists
107         if isinstance(x, tuple):
108             x = x[-1]
109 
110         if x.dim() == 4 and self.embedding_dim:
111             x = self.apply_pool(x)
112 
113         x = self.head(x)
114         return x
115 
116     def training_step(self, batch: Any, batch_idx: int) -> Any:
117         batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])
118         return super().training_step(batch, batch_idx)
119 
120     def validation_step(self, batch: Any, batch_idx: int) -> Any:
121         batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])
122         return super().validation_step(batch, batch_idx)
123 
124     def test_step(self, batch: Any, batch_idx: int) -> Any:
125         batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])
126         return super().test_step(batch, batch_idx)
127 
128     def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
129         batch = batch[DefaultDataKeys.INPUT]
130         return super().predict_step(batch, batch_idx, dataloader_idx=dataloader_idx)
131
```
--- 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/image/embedding/model.py b/flash/image/embedding/model.py
--- a/flash/image/embedding/model.py
+++ b/flash/image/embedding/model.py
@@ -107,7 +107,7 @@
         if isinstance(x, tuple):
             x = x[-1]
 
-        if x.dim() == 4 and self.embedding_dim:
+        if x.dim() == 4 and not self.embedding_dim:
             x = self.apply_pool(x)
 
         x = self.head(x)

{"golden_diff": "diff --git a/flash/image/embedding/model.py b/flash/image/embedding/model.py\n--- a/flash/image/embedding/model.py\n+++ b/flash/image/embedding/model.py\n@@ -107,7 +107,7 @@\n         if isinstance(x, tuple):\n             x = x[-1]\n \n-        if x.dim() == 4 and self.embedding_dim:\n+        if x.dim() == 4 and not self.embedding_dim:\n             x = self.apply_pool(x)\n \n         x = self.head(x)\n", "issue": "ImageEmbedder default behavior is not a flattened output\n## \ud83d\udc1b Bug\r\n\r\nI discovered this issue while testing PR #655. If you run the [Image Embedding README example code](https://github.com/PyTorchLightning/lightning-flash#example-1-image-embedding), it returns a 3D tensor. \r\nMy understanding from the use of embeddings in general, and how they are used in [Fifty One](https://voxel51.com/docs/fiftyone/tutorials/image_embeddings.html) is they expect the embeddings to be 1D (for each embedding). \r\n\r\nThe reason it returns a 3D tensor is because it depends on the backbone used. The default there is `resnet101`, which returns a `2048x7x7` shape tensor. Others like inception return a flat 1D tensor, i.e. length-X.\r\n\r\n### To Reproduce\r\n\r\nSteps to reproduce the behavior:\r\n\r\nRun the [README example](https://github.com/PyTorchLightning/lightning-flash#example-1-image-embedding), but remove the `embedding_dim` parameter. See below for example.\r\n\r\nNote: as-is, this will error on `print(embeddings.shape)`, regardless of configuration, since that is a list. But the question here is around the logic for the ImageEmbedder. \r\n\r\n\r\n#### Code sample\r\n```python\r\nfrom flash.core.data.utils import download_data\r\nfrom flash.image import ImageEmbedder\r\n\r\n# 1. Download the data\r\ndownload_data(\"https://pl-flash-data.s3.amazonaws.com/hymenoptera_data.zip\", \"data/\")\r\n\r\n# 2. Create an ImageEmbedder with resnet50 trained on imagenet.\r\nembedder = ImageEmbedder(backbone=\"resnet50\")\r\n\r\n# 3. Generate an embedding from an image path.\r\nembeddings = embedder.predict(\"data/hymenoptera_data/predict/153783656_85f9c3ac70.jpg\")\r\n\r\n# 4. Print embeddings shape\r\nprint(embeddings.shape)\r\n```\r\n\r\n### Expected behavior\r\n\r\nExpect to see a 100352x1 shape tensor as the output, instead of 2048x7x7. \r\n\r\n### Environment\r\n\r\n - PyTorch Version (e.g., 1.0): 1.9\r\n - OS (e.g., Linux): Linux\r\n - How you installed PyTorch (`conda`, `pip`, source): pip\r\n - Build command you used (if compiling from source): N/A\r\n - Python version: 3.8.6\r\n - CUDA/cuDNN version: N/A\r\n - GPU models and configuration: N/A\r\n - Any other relevant information: N/A\r\n\r\n### Additional context\r\n\r\nI believe the question is around what the logic should be here:\r\nhttps://github.com/PyTorchLightning/lightning-flash/blob/075de3a46d74d9fc0e769401063fede1f12d0518/flash/image/embedding/model.py#L85-L92\r\n\r\nIf `embedding_dim` is None, then the head is `nn.Identity()`. **If we desire a flat 1D embedding, then the question is: should `nn.Identity()` change to `nn.Flatten()`?**\r\n\r\nIt could be argued that the user should be left to flatten after on their own, but per the contributing guidelines, I thought this would align with \"[Force User Decisions To Best Practices](https://github.com/PyTorchLightning/lightning-flash/blob/ddd942d3dfe3884a97a855446410166c3c9f16d9/.github/CONTRIBUTING.md#force-user-decisions-to-best-practices)\"\r\n\r\nLet me know your thoughts. If that makes sense, then I can update the code, run some tests, and update docs in a PR. \r\n\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 Any, Callable, Mapping, Optional, Sequence, Tuple, Type, Union\n\nimport torch\nfrom pytorch_lightning.utilities import rank_zero_warn\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torchmetrics import Accuracy, Metric\n\nfrom flash.core.data.data_source import DefaultDataKeys\nfrom flash.core.model import Task\nfrom flash.core.registry import FlashRegistry\nfrom flash.core.utilities.imports import _IMAGE_AVAILABLE\nfrom flash.core.utilities.isinstance import _isinstance\nfrom flash.image.classification.data import ImageClassificationPreprocess\n\nif _IMAGE_AVAILABLE:\n    from flash.image.classification.backbones import IMAGE_CLASSIFIER_BACKBONES\nelse:\n    IMAGE_CLASSIFIER_BACKBONES = FlashRegistry(\"backbones\")\n\n\nclass ImageEmbedder(Task):\n    \"\"\"The ``ImageEmbedder`` is a :class:`~flash.Task` for obtaining feature vectors (embeddings) from images. For\n    more details, see :ref:`image_embedder`.\n\n    Args:\n        embedding_dim: Dimension of the embedded vector. ``None`` uses the default from the backbone.\n        backbone: A model to use to extract image features, defaults to ``\"swav-imagenet\"``.\n        pretrained: Use a pretrained backbone, defaults to ``True``.\n        loss_fn: Loss function for training and finetuning, defaults to :func:`torch.nn.functional.cross_entropy`\n        optimizer: Optimizer to use for training and finetuning, defaults to :class:`torch.optim.SGD`.\n        metrics: Metrics to compute for training and evaluation. Can either be an metric from the `torchmetrics`\n            package, a custom metric inherenting from `torchmetrics.Metric`, a callable function or a list/dict\n            containing a combination of the aforementioned. In all cases, each metric needs to have the signature\n            `metric(preds,target)` and return a single scalar tensor. Defaults to :class:`torchmetrics.Accuracy`.\n        learning_rate: Learning rate to use for training, defaults to ``1e-3``.\n        pooling_fn: Function used to pool image to generate embeddings, defaults to :func:`torch.max`.\n    \"\"\"\n\n    backbones: FlashRegistry = IMAGE_CLASSIFIER_BACKBONES\n\n    required_extras: str = \"image\"\n\n    def __init__(\n        self,\n        embedding_dim: Optional[int] = None,\n        backbone: str = \"resnet101\",\n        pretrained: bool = True,\n        loss_fn: Callable = F.cross_entropy,\n        optimizer: Type[torch.optim.Optimizer] = torch.optim.SGD,\n        metrics: Union[Metric, Callable, Mapping, Sequence, None] = (Accuracy()),\n        learning_rate: float = 1e-3,\n        pooling_fn: Callable = torch.max,\n    ):\n        super().__init__(\n            model=None,\n            loss_fn=loss_fn,\n            optimizer=optimizer,\n            metrics=metrics,\n            learning_rate=learning_rate,\n            preprocess=ImageClassificationPreprocess(),\n        )\n\n        self.save_hyperparameters()\n        self.backbone_name = backbone\n        self.embedding_dim = embedding_dim\n        assert pooling_fn in [torch.mean, torch.max]\n        self.pooling_fn = pooling_fn\n\n        self.backbone, num_features = self.backbones.get(backbone)(pretrained=pretrained)\n\n        if embedding_dim is None:\n            self.head = nn.Identity()\n        else:\n            self.head = nn.Sequential(\n                nn.Flatten(),\n                nn.Linear(num_features, embedding_dim),\n            )\n            rank_zero_warn(\"Adding linear layer on top of backbone. Remember to finetune first before using!\")\n\n    def apply_pool(self, x):\n        x = self.pooling_fn(x, dim=-1)\n        if _isinstance(x, Tuple[torch.Tensor, torch.Tensor]):\n            x = x[0]\n        x = self.pooling_fn(x, dim=-1)\n        if _isinstance(x, Tuple[torch.Tensor, torch.Tensor]):\n            x = x[0]\n        return x\n\n    def forward(self, x) -> torch.Tensor:\n        x = self.backbone(x)\n\n        # bolts ssl models return lists\n        if isinstance(x, tuple):\n            x = x[-1]\n\n        if x.dim() == 4 and self.embedding_dim:\n            x = self.apply_pool(x)\n\n        x = self.head(x)\n        return x\n\n    def training_step(self, batch: Any, batch_idx: int) -> Any:\n        batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])\n        return super().training_step(batch, batch_idx)\n\n    def validation_step(self, batch: Any, batch_idx: int) -> Any:\n        batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])\n        return super().validation_step(batch, batch_idx)\n\n    def test_step(self, batch: Any, batch_idx: int) -> Any:\n        batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])\n        return super().test_step(batch, batch_idx)\n\n    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:\n        batch = batch[DefaultDataKeys.INPUT]\n        return super().predict_step(batch, batch_idx, dataloader_idx=dataloader_idx)\n", "path": "flash/image/embedding/model.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 Any, Callable, Mapping, Optional, Sequence, Tuple, Type, Union\n\nimport torch\nfrom pytorch_lightning.utilities import rank_zero_warn\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torchmetrics import Accuracy, Metric\n\nfrom flash.core.data.data_source import DefaultDataKeys\nfrom flash.core.model import Task\nfrom flash.core.registry import FlashRegistry\nfrom flash.core.utilities.imports import _IMAGE_AVAILABLE\nfrom flash.core.utilities.isinstance import _isinstance\nfrom flash.image.classification.data import ImageClassificationPreprocess\n\nif _IMAGE_AVAILABLE:\n    from flash.image.classification.backbones import IMAGE_CLASSIFIER_BACKBONES\nelse:\n    IMAGE_CLASSIFIER_BACKBONES = FlashRegistry(\"backbones\")\n\n\nclass ImageEmbedder(Task):\n    \"\"\"The ``ImageEmbedder`` is a :class:`~flash.Task` for obtaining feature vectors (embeddings) from images. For\n    more details, see :ref:`image_embedder`.\n\n    Args:\n        embedding_dim: Dimension of the embedded vector. ``None`` uses the default from the backbone.\n        backbone: A model to use to extract image features, defaults to ``\"swav-imagenet\"``.\n        pretrained: Use a pretrained backbone, defaults to ``True``.\n        loss_fn: Loss function for training and finetuning, defaults to :func:`torch.nn.functional.cross_entropy`\n        optimizer: Optimizer to use for training and finetuning, defaults to :class:`torch.optim.SGD`.\n        metrics: Metrics to compute for training and evaluation. Can either be an metric from the `torchmetrics`\n            package, a custom metric inherenting from `torchmetrics.Metric`, a callable function or a list/dict\n            containing a combination of the aforementioned. In all cases, each metric needs to have the signature\n            `metric(preds,target)` and return a single scalar tensor. Defaults to :class:`torchmetrics.Accuracy`.\n        learning_rate: Learning rate to use for training, defaults to ``1e-3``.\n        pooling_fn: Function used to pool image to generate embeddings, defaults to :func:`torch.max`.\n    \"\"\"\n\n    backbones: FlashRegistry = IMAGE_CLASSIFIER_BACKBONES\n\n    required_extras: str = \"image\"\n\n    def __init__(\n        self,\n        embedding_dim: Optional[int] = None,\n        backbone: str = \"resnet101\",\n        pretrained: bool = True,\n        loss_fn: Callable = F.cross_entropy,\n        optimizer: Type[torch.optim.Optimizer] = torch.optim.SGD,\n        metrics: Union[Metric, Callable, Mapping, Sequence, None] = (Accuracy()),\n        learning_rate: float = 1e-3,\n        pooling_fn: Callable = torch.max,\n    ):\n        super().__init__(\n            model=None,\n            loss_fn=loss_fn,\n            optimizer=optimizer,\n            metrics=metrics,\n            learning_rate=learning_rate,\n            preprocess=ImageClassificationPreprocess(),\n        )\n\n        self.save_hyperparameters()\n        self.backbone_name = backbone\n        self.embedding_dim = embedding_dim\n        assert pooling_fn in [torch.mean, torch.max]\n        self.pooling_fn = pooling_fn\n\n        self.backbone, num_features = self.backbones.get(backbone)(pretrained=pretrained)\n\n        if embedding_dim is None:\n            self.head = nn.Identity()\n        else:\n            self.head = nn.Sequential(\n                nn.Flatten(),\n                nn.Linear(num_features, embedding_dim),\n            )\n            rank_zero_warn(\"Adding linear layer on top of backbone. Remember to finetune first before using!\")\n\n    def apply_pool(self, x):\n        x = self.pooling_fn(x, dim=-1)\n        if _isinstance(x, Tuple[torch.Tensor, torch.Tensor]):\n            x = x[0]\n        x = self.pooling_fn(x, dim=-1)\n        if _isinstance(x, Tuple[torch.Tensor, torch.Tensor]):\n            x = x[0]\n        return x\n\n    def forward(self, x) -> torch.Tensor:\n        x = self.backbone(x)\n\n        # bolts ssl models return lists\n        if isinstance(x, tuple):\n            x = x[-1]\n\n        if x.dim() == 4 and not self.embedding_dim:\n            x = self.apply_pool(x)\n\n        x = self.head(x)\n        return x\n\n    def training_step(self, batch: Any, batch_idx: int) -> Any:\n        batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])\n        return super().training_step(batch, batch_idx)\n\n    def validation_step(self, batch: Any, batch_idx: int) -> Any:\n        batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])\n        return super().validation_step(batch, batch_idx)\n\n    def test_step(self, batch: Any, batch_idx: int) -> Any:\n        batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])\n        return super().test_step(batch, batch_idx)\n\n    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:\n        batch = batch[DefaultDataKeys.INPUT]\n        return super().predict_step(batch, batch_idx, dataloader_idx=dataloader_idx)\n", "path": "flash/image/embedding/model.py"}]}