We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
`METADATA` can't be a dict
## 🐛 Bug

If a dict is used for the `METADATA` key then it will give an error. This should be supported.


--- 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/core/data/batch.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, Dict, List, Mapping, Optional, Sequence, Tuple, TYPE_CHECKING, Union
15 
16 import torch
17 from pytorch_lightning.trainer.states import RunningStage
18 from pytorch_lightning.utilities.exceptions import MisconfigurationException
19 from torch import Tensor
20 
21 from flash.core.data.callback import ControlFlow
22 from flash.core.data.data_source import DefaultDataKeys
23 from flash.core.data.utils import (
24     _contains_any_tensor,
25     convert_to_modules,
26     CurrentFuncContext,
27     CurrentRunningStageContext,
28 )
29 
30 if TYPE_CHECKING:  # pragma: no-cover
31     from flash.core.data.process import Preprocess
32 
33 
34 class _Sequential(torch.nn.Module):
35     """
36     This class is used to chain 3 functions together for the _Preprocessor ``per_sample_transform`` function.
37     1. ``pre_tensor_transform``
38     2. ``to_tensor_transform``
39     3. ``post_tensor_transform``
40     """
41 
42     def __init__(
43         self,
44         preprocess: 'Preprocess',
45         pre_tensor_transform: Callable,
46         to_tensor_transform: Callable,
47         post_tensor_transform: Callable,
48         stage: RunningStage,
49         assert_contains_tensor: bool = False,
50     ):
51         super().__init__()
52         self.preprocess = preprocess
53         self.callback = ControlFlow(self.preprocess.callbacks)
54         self.pre_tensor_transform = convert_to_modules(pre_tensor_transform)
55         self.to_tensor_transform = convert_to_modules(to_tensor_transform)
56         self.post_tensor_transform = convert_to_modules(post_tensor_transform)
57         self.stage = stage
58         self.assert_contains_tensor = assert_contains_tensor
59 
60         self._current_stage_context = CurrentRunningStageContext(stage, preprocess, reset=False)
61         self._pre_tensor_transform_context = CurrentFuncContext("pre_tensor_transform", preprocess)
62         self._to_tensor_transform_context = CurrentFuncContext("to_tensor_transform", preprocess)
63         self._post_tensor_transform_context = CurrentFuncContext("post_tensor_transform", preprocess)
64 
65     def forward(self, sample: Any) -> Any:
66         self.callback.on_load_sample(sample, self.stage)
67 
68         with self._current_stage_context:
69             with self._pre_tensor_transform_context:
70                 sample = self.pre_tensor_transform(sample)
71                 self.callback.on_pre_tensor_transform(sample, self.stage)
72 
73             with self._to_tensor_transform_context:
74                 sample = self.to_tensor_transform(sample)
75                 self.callback.on_to_tensor_transform(sample, self.stage)
76 
77             if self.assert_contains_tensor:
78                 if not _contains_any_tensor(sample):
79                     raise MisconfigurationException(
80                         "When ``to_tensor_transform`` is overriden, "
81                         "``DataPipeline`` expects the outputs to be ``tensors``"
82                     )
83 
84             with self._post_tensor_transform_context:
85                 sample = self.post_tensor_transform(sample)
86                 self.callback.on_post_tensor_transform(sample, self.stage)
87 
88             return sample
89 
90     def __str__(self) -> str:
91         return (
92             f"{self.__class__.__name__}:\n"
93             f"\t(pre_tensor_transform): {str(self.pre_tensor_transform)}\n"
94             f"\t(to_tensor_transform): {str(self.to_tensor_transform)}\n"
95             f"\t(post_tensor_transform): {str(self.post_tensor_transform)}\n"
96             f"\t(assert_contains_tensor): {str(self.assert_contains_tensor)}\n"
97             f"\t(stage): {str(self.stage)}"
98         )
99 
100 
101 class _Preprocessor(torch.nn.Module):
102     """
103         This class is used to encapsultate the following functions of a Preprocess Object:
104         Inside a worker:
105             per_sample_transform: Function to transform an individual sample
106                 Inside a worker, it is actually make of 3 functions:
107                     * pre_tensor_transform
108                     * to_tensor_transform
109                     * post_tensor_transform
110             collate: Function to merge sample into a batch
111             per_batch_transform: Function to transform an individual batch
112                 * per_batch_transform
113 
114         Inside main process:
115             per_sample_transform: Function to transform an individual sample
116                 * per_sample_transform_on_device
117             collate: Function to merge sample into a batch
118             per_batch_transform: Function to transform an individual batch
119                 * per_batch_transform_on_device
120     """
121 
122     def __init__(
123         self,
124         preprocess: 'Preprocess',
125         collate_fn: Callable,
126         per_sample_transform: Union[Callable, _Sequential],
127         per_batch_transform: Callable,
128         stage: RunningStage,
129         apply_per_sample_transform: bool = True,
130         on_device: bool = False,
131     ):
132         super().__init__()
133         self.preprocess = preprocess
134         self.callback = ControlFlow(self.preprocess.callbacks)
135         self.collate_fn = convert_to_modules(collate_fn)
136         self.per_sample_transform = convert_to_modules(per_sample_transform)
137         self.per_batch_transform = convert_to_modules(per_batch_transform)
138         self.apply_per_sample_transform = apply_per_sample_transform
139         self.stage = stage
140         self.on_device = on_device
141 
142         extension = f"{'_on_device' if self.on_device else ''}"
143         self._current_stage_context = CurrentRunningStageContext(stage, preprocess)
144         self._per_sample_transform_context = CurrentFuncContext(f"per_sample_transform{extension}", preprocess)
145         self._collate_context = CurrentFuncContext("collate", preprocess)
146         self._per_batch_transform_context = CurrentFuncContext(f"per_batch_transform{extension}", preprocess)
147 
148     def _extract_metadata(
149         self,
150         samples: List[Dict[str, Any]],
151     ) -> Tuple[List[Dict[str, Any]], Optional[List[Dict[str, Any]]]]:
152         metadata = [s.pop(DefaultDataKeys.METADATA, None) if isinstance(s, Mapping) else None for s in samples]
153         return samples, metadata if any(m is not None for m in metadata) else None
154 
155     def forward(self, samples: Sequence[Any]) -> Any:
156         # we create a new dict to prevent from potential memory leaks
157         # assuming that the dictionary samples are stored in between and
158         # potentially modified before the transforms are applied.
159         if isinstance(samples, dict):
160             samples = dict(samples.items())
161 
162         with self._current_stage_context:
163 
164             if self.apply_per_sample_transform:
165                 with self._per_sample_transform_context:
166                     _samples = []
167                     for sample in samples:
168                         sample = self.per_sample_transform(sample)
169                         if self.on_device:
170                             self.callback.on_per_sample_transform_on_device(sample, self.stage)
171                         _samples.append(sample)
172 
173                 samples = type(_samples)(_samples)
174 
175                 with self._collate_context:
176                     samples, metadata = self._extract_metadata(samples)
177                     samples = self.collate_fn(samples)
178                     if metadata:
179                         samples[DefaultDataKeys.METADATA] = metadata
180                     self.callback.on_collate(samples, self.stage)
181 
182             with self._per_batch_transform_context:
183                 samples = self.per_batch_transform(samples)
184                 if self.on_device:
185                     self.callback.on_per_batch_transform_on_device(samples, self.stage)
186                 else:
187                     self.callback.on_per_batch_transform(samples, self.stage)
188             return samples
189 
190     def __str__(self) -> str:
191         # todo: define repr function which would take object and string attributes to be shown
192         return (
193             "_Preprocessor:\n"
194             f"\t(per_sample_transform): {str(self.per_sample_transform)}\n"
195             f"\t(collate_fn): {str(self.collate_fn)}\n"
196             f"\t(per_batch_transform): {str(self.per_batch_transform)}\n"
197             f"\t(apply_per_sample_transform): {str(self.apply_per_sample_transform)}\n"
198             f"\t(on_device): {str(self.on_device)}\n"
199             f"\t(stage): {str(self.stage)}"
200         )
201 
202 
203 class _Postprocessor(torch.nn.Module):
204     """
205         This class is used to encapsultate the following functions of a Postprocess Object:
206         Inside main process:
207             per_batch_transform: Function to transform a batch
208             per_sample_transform: Function to transform an individual sample
209             uncollate_fn: Function to split a batch into samples
210             per_sample_transform: Function to transform an individual sample
211             save_fn: Function to save all data
212             save_per_sample: Function to save an individual sample
213     """
214 
215     def __init__(
216         self,
217         uncollate_fn: Callable,
218         per_batch_transform: Callable,
219         per_sample_transform: Callable,
220         serializer: Optional[Callable],
221         save_fn: Optional[Callable] = None,
222         save_per_sample: bool = False
223     ):
224         super().__init__()
225         self.uncollate_fn = convert_to_modules(uncollate_fn)
226         self.per_batch_transform = convert_to_modules(per_batch_transform)
227         self.per_sample_transform = convert_to_modules(per_sample_transform)
228         self.serializer = convert_to_modules(serializer)
229         self.save_fn = convert_to_modules(save_fn)
230         self.save_per_sample = convert_to_modules(save_per_sample)
231 
232     def forward(self, batch: Sequence[Any]):
233         uncollated = self.uncollate_fn(self.per_batch_transform(batch))
234 
235         final_preds = type(uncollated)([self.serializer(self.per_sample_transform(sample)) for sample in uncollated])
236 
237         if self.save_fn:
238             if self.save_per_sample:
239                 for pred in final_preds:
240                     self.save_fn(pred)
241             else:
242                 self.save_fn(final_preds)
243         else:
244             return final_preds
245 
246     def __str__(self) -> str:
247         return (
248             "_Postprocessor:\n"
249             f"\t(per_batch_transform): {str(self.per_batch_transform)}\n"
250             f"\t(uncollate_fn): {str(self.uncollate_fn)}\n"
251             f"\t(per_sample_transform): {str(self.per_sample_transform)}\n"
252             f"\t(serializer): {str(self.serializer)}"
253         )
254 
255 
256 def default_uncollate(batch: Any):
257     """
258     This function is used to uncollate a batch into samples.
259 
260     Examples:
261         >>> a, b = default_uncollate(torch.rand((2,1)))
262     """
263 
264     batch_type = type(batch)
265 
266     if isinstance(batch, Tensor):
267         if len(batch.shape) == 0:  # 0 shape tensors
268             return batch
269         return list(torch.unbind(batch, 0))
270 
271     elif isinstance(batch, Mapping):
272         return [batch_type(dict(zip(batch, default_uncollate(t)))) for t in zip(*batch.values())]
273 
274     elif isinstance(batch, tuple) and hasattr(batch, '_fields'):  # namedtuple
275         return [batch_type(*default_uncollate(sample)) for sample in zip(*batch)]
276 
277     elif isinstance(batch, Sequence) and not isinstance(batch, str):
278         return [default_uncollate(sample) for sample in batch]
279 
280     return batch
281
```
--- 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/core/data/batch.py b/flash/core/data/batch.py
--- a/flash/core/data/batch.py
+++ b/flash/core/data/batch.py
@@ -145,8 +145,8 @@
         self._collate_context = CurrentFuncContext("collate", preprocess)
         self._per_batch_transform_context = CurrentFuncContext(f"per_batch_transform{extension}", preprocess)
 
+    @staticmethod
     def _extract_metadata(
-        self,
         samples: List[Dict[str, Any]],
     ) -> Tuple[List[Dict[str, Any]], Optional[List[Dict[str, Any]]]]:
         metadata = [s.pop(DefaultDataKeys.METADATA, None) if isinstance(s, Mapping) else None for s in samples]
@@ -229,8 +229,18 @@
         self.save_fn = convert_to_modules(save_fn)
         self.save_per_sample = convert_to_modules(save_per_sample)
 
+    @staticmethod
+    def _extract_metadata(batch: Any) -> Tuple[Any, Optional[Any]]:
+        if isinstance(batch, Mapping):
+            return batch, batch.get(DefaultDataKeys.METADATA, None)
+        return batch, None
+
     def forward(self, batch: Sequence[Any]):
+        batch, metadata = self._extract_metadata(batch)
         uncollated = self.uncollate_fn(self.per_batch_transform(batch))
+        if metadata:
+            for sample, sample_metadata in zip(uncollated, metadata):
+                sample[DefaultDataKeys.METADATA] = sample_metadata
 
         final_preds = type(uncollated)([self.serializer(self.per_sample_transform(sample)) for sample in uncollated])

{"golden_diff": "diff --git a/flash/core/data/batch.py b/flash/core/data/batch.py\n--- a/flash/core/data/batch.py\n+++ b/flash/core/data/batch.py\n@@ -145,8 +145,8 @@\n         self._collate_context = CurrentFuncContext(\"collate\", preprocess)\n         self._per_batch_transform_context = CurrentFuncContext(f\"per_batch_transform{extension}\", preprocess)\n \n+    @staticmethod\n     def _extract_metadata(\n-        self,\n         samples: List[Dict[str, Any]],\n     ) -> Tuple[List[Dict[str, Any]], Optional[List[Dict[str, Any]]]]:\n         metadata = [s.pop(DefaultDataKeys.METADATA, None) if isinstance(s, Mapping) else None for s in samples]\n@@ -229,8 +229,18 @@\n         self.save_fn = convert_to_modules(save_fn)\n         self.save_per_sample = convert_to_modules(save_per_sample)\n \n+    @staticmethod\n+    def _extract_metadata(batch: Any) -> Tuple[Any, Optional[Any]]:\n+        if isinstance(batch, Mapping):\n+            return batch, batch.get(DefaultDataKeys.METADATA, None)\n+        return batch, None\n+\n     def forward(self, batch: Sequence[Any]):\n+        batch, metadata = self._extract_metadata(batch)\n         uncollated = self.uncollate_fn(self.per_batch_transform(batch))\n+        if metadata:\n+            for sample, sample_metadata in zip(uncollated, metadata):\n+                sample[DefaultDataKeys.METADATA] = sample_metadata\n \n         final_preds = type(uncollated)([self.serializer(self.per_sample_transform(sample)) for sample in uncollated])\n", "issue": "`METADATA` can't be a dict\n## \ud83d\udc1b Bug\r\n\r\nIf a dict is used for the `METADATA` key then it will give an error. This should be supported.\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, Dict, List, Mapping, Optional, Sequence, Tuple, TYPE_CHECKING, Union\n\nimport torch\nfrom pytorch_lightning.trainer.states import RunningStage\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\nfrom torch import Tensor\n\nfrom flash.core.data.callback import ControlFlow\nfrom flash.core.data.data_source import DefaultDataKeys\nfrom flash.core.data.utils import (\n    _contains_any_tensor,\n    convert_to_modules,\n    CurrentFuncContext,\n    CurrentRunningStageContext,\n)\n\nif TYPE_CHECKING:  # pragma: no-cover\n    from flash.core.data.process import Preprocess\n\n\nclass _Sequential(torch.nn.Module):\n    \"\"\"\n    This class is used to chain 3 functions together for the _Preprocessor ``per_sample_transform`` function.\n    1. ``pre_tensor_transform``\n    2. ``to_tensor_transform``\n    3. ``post_tensor_transform``\n    \"\"\"\n\n    def __init__(\n        self,\n        preprocess: 'Preprocess',\n        pre_tensor_transform: Callable,\n        to_tensor_transform: Callable,\n        post_tensor_transform: Callable,\n        stage: RunningStage,\n        assert_contains_tensor: bool = False,\n    ):\n        super().__init__()\n        self.preprocess = preprocess\n        self.callback = ControlFlow(self.preprocess.callbacks)\n        self.pre_tensor_transform = convert_to_modules(pre_tensor_transform)\n        self.to_tensor_transform = convert_to_modules(to_tensor_transform)\n        self.post_tensor_transform = convert_to_modules(post_tensor_transform)\n        self.stage = stage\n        self.assert_contains_tensor = assert_contains_tensor\n\n        self._current_stage_context = CurrentRunningStageContext(stage, preprocess, reset=False)\n        self._pre_tensor_transform_context = CurrentFuncContext(\"pre_tensor_transform\", preprocess)\n        self._to_tensor_transform_context = CurrentFuncContext(\"to_tensor_transform\", preprocess)\n        self._post_tensor_transform_context = CurrentFuncContext(\"post_tensor_transform\", preprocess)\n\n    def forward(self, sample: Any) -> Any:\n        self.callback.on_load_sample(sample, self.stage)\n\n        with self._current_stage_context:\n            with self._pre_tensor_transform_context:\n                sample = self.pre_tensor_transform(sample)\n                self.callback.on_pre_tensor_transform(sample, self.stage)\n\n            with self._to_tensor_transform_context:\n                sample = self.to_tensor_transform(sample)\n                self.callback.on_to_tensor_transform(sample, self.stage)\n\n            if self.assert_contains_tensor:\n                if not _contains_any_tensor(sample):\n                    raise MisconfigurationException(\n                        \"When ``to_tensor_transform`` is overriden, \"\n                        \"``DataPipeline`` expects the outputs to be ``tensors``\"\n                    )\n\n            with self._post_tensor_transform_context:\n                sample = self.post_tensor_transform(sample)\n                self.callback.on_post_tensor_transform(sample, self.stage)\n\n            return sample\n\n    def __str__(self) -> str:\n        return (\n            f\"{self.__class__.__name__}:\\n\"\n            f\"\\t(pre_tensor_transform): {str(self.pre_tensor_transform)}\\n\"\n            f\"\\t(to_tensor_transform): {str(self.to_tensor_transform)}\\n\"\n            f\"\\t(post_tensor_transform): {str(self.post_tensor_transform)}\\n\"\n            f\"\\t(assert_contains_tensor): {str(self.assert_contains_tensor)}\\n\"\n            f\"\\t(stage): {str(self.stage)}\"\n        )\n\n\nclass _Preprocessor(torch.nn.Module):\n    \"\"\"\n        This class is used to encapsultate the following functions of a Preprocess Object:\n        Inside a worker:\n            per_sample_transform: Function to transform an individual sample\n                Inside a worker, it is actually make of 3 functions:\n                    * pre_tensor_transform\n                    * to_tensor_transform\n                    * post_tensor_transform\n            collate: Function to merge sample into a batch\n            per_batch_transform: Function to transform an individual batch\n                * per_batch_transform\n\n        Inside main process:\n            per_sample_transform: Function to transform an individual sample\n                * per_sample_transform_on_device\n            collate: Function to merge sample into a batch\n            per_batch_transform: Function to transform an individual batch\n                * per_batch_transform_on_device\n    \"\"\"\n\n    def __init__(\n        self,\n        preprocess: 'Preprocess',\n        collate_fn: Callable,\n        per_sample_transform: Union[Callable, _Sequential],\n        per_batch_transform: Callable,\n        stage: RunningStage,\n        apply_per_sample_transform: bool = True,\n        on_device: bool = False,\n    ):\n        super().__init__()\n        self.preprocess = preprocess\n        self.callback = ControlFlow(self.preprocess.callbacks)\n        self.collate_fn = convert_to_modules(collate_fn)\n        self.per_sample_transform = convert_to_modules(per_sample_transform)\n        self.per_batch_transform = convert_to_modules(per_batch_transform)\n        self.apply_per_sample_transform = apply_per_sample_transform\n        self.stage = stage\n        self.on_device = on_device\n\n        extension = f\"{'_on_device' if self.on_device else ''}\"\n        self._current_stage_context = CurrentRunningStageContext(stage, preprocess)\n        self._per_sample_transform_context = CurrentFuncContext(f\"per_sample_transform{extension}\", preprocess)\n        self._collate_context = CurrentFuncContext(\"collate\", preprocess)\n        self._per_batch_transform_context = CurrentFuncContext(f\"per_batch_transform{extension}\", preprocess)\n\n    def _extract_metadata(\n        self,\n        samples: List[Dict[str, Any]],\n    ) -> Tuple[List[Dict[str, Any]], Optional[List[Dict[str, Any]]]]:\n        metadata = [s.pop(DefaultDataKeys.METADATA, None) if isinstance(s, Mapping) else None for s in samples]\n        return samples, metadata if any(m is not None for m in metadata) else None\n\n    def forward(self, samples: Sequence[Any]) -> Any:\n        # we create a new dict to prevent from potential memory leaks\n        # assuming that the dictionary samples are stored in between and\n        # potentially modified before the transforms are applied.\n        if isinstance(samples, dict):\n            samples = dict(samples.items())\n\n        with self._current_stage_context:\n\n            if self.apply_per_sample_transform:\n                with self._per_sample_transform_context:\n                    _samples = []\n                    for sample in samples:\n                        sample = self.per_sample_transform(sample)\n                        if self.on_device:\n                            self.callback.on_per_sample_transform_on_device(sample, self.stage)\n                        _samples.append(sample)\n\n                samples = type(_samples)(_samples)\n\n                with self._collate_context:\n                    samples, metadata = self._extract_metadata(samples)\n                    samples = self.collate_fn(samples)\n                    if metadata:\n                        samples[DefaultDataKeys.METADATA] = metadata\n                    self.callback.on_collate(samples, self.stage)\n\n            with self._per_batch_transform_context:\n                samples = self.per_batch_transform(samples)\n                if self.on_device:\n                    self.callback.on_per_batch_transform_on_device(samples, self.stage)\n                else:\n                    self.callback.on_per_batch_transform(samples, self.stage)\n            return samples\n\n    def __str__(self) -> str:\n        # todo: define repr function which would take object and string attributes to be shown\n        return (\n            \"_Preprocessor:\\n\"\n            f\"\\t(per_sample_transform): {str(self.per_sample_transform)}\\n\"\n            f\"\\t(collate_fn): {str(self.collate_fn)}\\n\"\n            f\"\\t(per_batch_transform): {str(self.per_batch_transform)}\\n\"\n            f\"\\t(apply_per_sample_transform): {str(self.apply_per_sample_transform)}\\n\"\n            f\"\\t(on_device): {str(self.on_device)}\\n\"\n            f\"\\t(stage): {str(self.stage)}\"\n        )\n\n\nclass _Postprocessor(torch.nn.Module):\n    \"\"\"\n        This class is used to encapsultate the following functions of a Postprocess Object:\n        Inside main process:\n            per_batch_transform: Function to transform a batch\n            per_sample_transform: Function to transform an individual sample\n            uncollate_fn: Function to split a batch into samples\n            per_sample_transform: Function to transform an individual sample\n            save_fn: Function to save all data\n            save_per_sample: Function to save an individual sample\n    \"\"\"\n\n    def __init__(\n        self,\n        uncollate_fn: Callable,\n        per_batch_transform: Callable,\n        per_sample_transform: Callable,\n        serializer: Optional[Callable],\n        save_fn: Optional[Callable] = None,\n        save_per_sample: bool = False\n    ):\n        super().__init__()\n        self.uncollate_fn = convert_to_modules(uncollate_fn)\n        self.per_batch_transform = convert_to_modules(per_batch_transform)\n        self.per_sample_transform = convert_to_modules(per_sample_transform)\n        self.serializer = convert_to_modules(serializer)\n        self.save_fn = convert_to_modules(save_fn)\n        self.save_per_sample = convert_to_modules(save_per_sample)\n\n    def forward(self, batch: Sequence[Any]):\n        uncollated = self.uncollate_fn(self.per_batch_transform(batch))\n\n        final_preds = type(uncollated)([self.serializer(self.per_sample_transform(sample)) for sample in uncollated])\n\n        if self.save_fn:\n            if self.save_per_sample:\n                for pred in final_preds:\n                    self.save_fn(pred)\n            else:\n                self.save_fn(final_preds)\n        else:\n            return final_preds\n\n    def __str__(self) -> str:\n        return (\n            \"_Postprocessor:\\n\"\n            f\"\\t(per_batch_transform): {str(self.per_batch_transform)}\\n\"\n            f\"\\t(uncollate_fn): {str(self.uncollate_fn)}\\n\"\n            f\"\\t(per_sample_transform): {str(self.per_sample_transform)}\\n\"\n            f\"\\t(serializer): {str(self.serializer)}\"\n        )\n\n\ndef default_uncollate(batch: Any):\n    \"\"\"\n    This function is used to uncollate a batch into samples.\n\n    Examples:\n        >>> a, b = default_uncollate(torch.rand((2,1)))\n    \"\"\"\n\n    batch_type = type(batch)\n\n    if isinstance(batch, Tensor):\n        if len(batch.shape) == 0:  # 0 shape tensors\n            return batch\n        return list(torch.unbind(batch, 0))\n\n    elif isinstance(batch, Mapping):\n        return [batch_type(dict(zip(batch, default_uncollate(t)))) for t in zip(*batch.values())]\n\n    elif isinstance(batch, tuple) and hasattr(batch, '_fields'):  # namedtuple\n        return [batch_type(*default_uncollate(sample)) for sample in zip(*batch)]\n\n    elif isinstance(batch, Sequence) and not isinstance(batch, str):\n        return [default_uncollate(sample) for sample in batch]\n\n    return batch\n", "path": "flash/core/data/batch.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, Dict, List, Mapping, Optional, Sequence, Tuple, TYPE_CHECKING, Union\n\nimport torch\nfrom pytorch_lightning.trainer.states import RunningStage\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\nfrom torch import Tensor\n\nfrom flash.core.data.callback import ControlFlow\nfrom flash.core.data.data_source import DefaultDataKeys\nfrom flash.core.data.utils import (\n    _contains_any_tensor,\n    convert_to_modules,\n    CurrentFuncContext,\n    CurrentRunningStageContext,\n)\n\nif TYPE_CHECKING:  # pragma: no-cover\n    from flash.core.data.process import Preprocess\n\n\nclass _Sequential(torch.nn.Module):\n    \"\"\"\n    This class is used to chain 3 functions together for the _Preprocessor ``per_sample_transform`` function.\n    1. ``pre_tensor_transform``\n    2. ``to_tensor_transform``\n    3. ``post_tensor_transform``\n    \"\"\"\n\n    def __init__(\n        self,\n        preprocess: 'Preprocess',\n        pre_tensor_transform: Callable,\n        to_tensor_transform: Callable,\n        post_tensor_transform: Callable,\n        stage: RunningStage,\n        assert_contains_tensor: bool = False,\n    ):\n        super().__init__()\n        self.preprocess = preprocess\n        self.callback = ControlFlow(self.preprocess.callbacks)\n        self.pre_tensor_transform = convert_to_modules(pre_tensor_transform)\n        self.to_tensor_transform = convert_to_modules(to_tensor_transform)\n        self.post_tensor_transform = convert_to_modules(post_tensor_transform)\n        self.stage = stage\n        self.assert_contains_tensor = assert_contains_tensor\n\n        self._current_stage_context = CurrentRunningStageContext(stage, preprocess, reset=False)\n        self._pre_tensor_transform_context = CurrentFuncContext(\"pre_tensor_transform\", preprocess)\n        self._to_tensor_transform_context = CurrentFuncContext(\"to_tensor_transform\", preprocess)\n        self._post_tensor_transform_context = CurrentFuncContext(\"post_tensor_transform\", preprocess)\n\n    def forward(self, sample: Any) -> Any:\n        self.callback.on_load_sample(sample, self.stage)\n\n        with self._current_stage_context:\n            with self._pre_tensor_transform_context:\n                sample = self.pre_tensor_transform(sample)\n                self.callback.on_pre_tensor_transform(sample, self.stage)\n\n            with self._to_tensor_transform_context:\n                sample = self.to_tensor_transform(sample)\n                self.callback.on_to_tensor_transform(sample, self.stage)\n\n            if self.assert_contains_tensor:\n                if not _contains_any_tensor(sample):\n                    raise MisconfigurationException(\n                        \"When ``to_tensor_transform`` is overriden, \"\n                        \"``DataPipeline`` expects the outputs to be ``tensors``\"\n                    )\n\n            with self._post_tensor_transform_context:\n                sample = self.post_tensor_transform(sample)\n                self.callback.on_post_tensor_transform(sample, self.stage)\n\n            return sample\n\n    def __str__(self) -> str:\n        return (\n            f\"{self.__class__.__name__}:\\n\"\n            f\"\\t(pre_tensor_transform): {str(self.pre_tensor_transform)}\\n\"\n            f\"\\t(to_tensor_transform): {str(self.to_tensor_transform)}\\n\"\n            f\"\\t(post_tensor_transform): {str(self.post_tensor_transform)}\\n\"\n            f\"\\t(assert_contains_tensor): {str(self.assert_contains_tensor)}\\n\"\n            f\"\\t(stage): {str(self.stage)}\"\n        )\n\n\nclass _Preprocessor(torch.nn.Module):\n    \"\"\"\n        This class is used to encapsultate the following functions of a Preprocess Object:\n        Inside a worker:\n            per_sample_transform: Function to transform an individual sample\n                Inside a worker, it is actually make of 3 functions:\n                    * pre_tensor_transform\n                    * to_tensor_transform\n                    * post_tensor_transform\n            collate: Function to merge sample into a batch\n            per_batch_transform: Function to transform an individual batch\n                * per_batch_transform\n\n        Inside main process:\n            per_sample_transform: Function to transform an individual sample\n                * per_sample_transform_on_device\n            collate: Function to merge sample into a batch\n            per_batch_transform: Function to transform an individual batch\n                * per_batch_transform_on_device\n    \"\"\"\n\n    def __init__(\n        self,\n        preprocess: 'Preprocess',\n        collate_fn: Callable,\n        per_sample_transform: Union[Callable, _Sequential],\n        per_batch_transform: Callable,\n        stage: RunningStage,\n        apply_per_sample_transform: bool = True,\n        on_device: bool = False,\n    ):\n        super().__init__()\n        self.preprocess = preprocess\n        self.callback = ControlFlow(self.preprocess.callbacks)\n        self.collate_fn = convert_to_modules(collate_fn)\n        self.per_sample_transform = convert_to_modules(per_sample_transform)\n        self.per_batch_transform = convert_to_modules(per_batch_transform)\n        self.apply_per_sample_transform = apply_per_sample_transform\n        self.stage = stage\n        self.on_device = on_device\n\n        extension = f\"{'_on_device' if self.on_device else ''}\"\n        self._current_stage_context = CurrentRunningStageContext(stage, preprocess)\n        self._per_sample_transform_context = CurrentFuncContext(f\"per_sample_transform{extension}\", preprocess)\n        self._collate_context = CurrentFuncContext(\"collate\", preprocess)\n        self._per_batch_transform_context = CurrentFuncContext(f\"per_batch_transform{extension}\", preprocess)\n\n    @staticmethod\n    def _extract_metadata(\n        samples: List[Dict[str, Any]],\n    ) -> Tuple[List[Dict[str, Any]], Optional[List[Dict[str, Any]]]]:\n        metadata = [s.pop(DefaultDataKeys.METADATA, None) if isinstance(s, Mapping) else None for s in samples]\n        return samples, metadata if any(m is not None for m in metadata) else None\n\n    def forward(self, samples: Sequence[Any]) -> Any:\n        # we create a new dict to prevent from potential memory leaks\n        # assuming that the dictionary samples are stored in between and\n        # potentially modified before the transforms are applied.\n        if isinstance(samples, dict):\n            samples = dict(samples.items())\n\n        with self._current_stage_context:\n\n            if self.apply_per_sample_transform:\n                with self._per_sample_transform_context:\n                    _samples = []\n                    for sample in samples:\n                        sample = self.per_sample_transform(sample)\n                        if self.on_device:\n                            self.callback.on_per_sample_transform_on_device(sample, self.stage)\n                        _samples.append(sample)\n\n                samples = type(_samples)(_samples)\n\n                with self._collate_context:\n                    samples, metadata = self._extract_metadata(samples)\n                    samples = self.collate_fn(samples)\n                    if metadata:\n                        samples[DefaultDataKeys.METADATA] = metadata\n                    self.callback.on_collate(samples, self.stage)\n\n            with self._per_batch_transform_context:\n                samples = self.per_batch_transform(samples)\n                if self.on_device:\n                    self.callback.on_per_batch_transform_on_device(samples, self.stage)\n                else:\n                    self.callback.on_per_batch_transform(samples, self.stage)\n            return samples\n\n    def __str__(self) -> str:\n        # todo: define repr function which would take object and string attributes to be shown\n        return (\n            \"_Preprocessor:\\n\"\n            f\"\\t(per_sample_transform): {str(self.per_sample_transform)}\\n\"\n            f\"\\t(collate_fn): {str(self.collate_fn)}\\n\"\n            f\"\\t(per_batch_transform): {str(self.per_batch_transform)}\\n\"\n            f\"\\t(apply_per_sample_transform): {str(self.apply_per_sample_transform)}\\n\"\n            f\"\\t(on_device): {str(self.on_device)}\\n\"\n            f\"\\t(stage): {str(self.stage)}\"\n        )\n\n\nclass _Postprocessor(torch.nn.Module):\n    \"\"\"\n        This class is used to encapsultate the following functions of a Postprocess Object:\n        Inside main process:\n            per_batch_transform: Function to transform a batch\n            per_sample_transform: Function to transform an individual sample\n            uncollate_fn: Function to split a batch into samples\n            per_sample_transform: Function to transform an individual sample\n            save_fn: Function to save all data\n            save_per_sample: Function to save an individual sample\n    \"\"\"\n\n    def __init__(\n        self,\n        uncollate_fn: Callable,\n        per_batch_transform: Callable,\n        per_sample_transform: Callable,\n        serializer: Optional[Callable],\n        save_fn: Optional[Callable] = None,\n        save_per_sample: bool = False\n    ):\n        super().__init__()\n        self.uncollate_fn = convert_to_modules(uncollate_fn)\n        self.per_batch_transform = convert_to_modules(per_batch_transform)\n        self.per_sample_transform = convert_to_modules(per_sample_transform)\n        self.serializer = convert_to_modules(serializer)\n        self.save_fn = convert_to_modules(save_fn)\n        self.save_per_sample = convert_to_modules(save_per_sample)\n\n    @staticmethod\n    def _extract_metadata(batch: Any) -> Tuple[Any, Optional[Any]]:\n        if isinstance(batch, Mapping):\n            return batch, batch.get(DefaultDataKeys.METADATA, None)\n        return batch, None\n\n    def forward(self, batch: Sequence[Any]):\n        batch, metadata = self._extract_metadata(batch)\n        uncollated = self.uncollate_fn(self.per_batch_transform(batch))\n        if metadata:\n            for sample, sample_metadata in zip(uncollated, metadata):\n                sample[DefaultDataKeys.METADATA] = sample_metadata\n\n        final_preds = type(uncollated)([self.serializer(self.per_sample_transform(sample)) for sample in uncollated])\n\n        if self.save_fn:\n            if self.save_per_sample:\n                for pred in final_preds:\n                    self.save_fn(pred)\n            else:\n                self.save_fn(final_preds)\n        else:\n            return final_preds\n\n    def __str__(self) -> str:\n        return (\n            \"_Postprocessor:\\n\"\n            f\"\\t(per_batch_transform): {str(self.per_batch_transform)}\\n\"\n            f\"\\t(uncollate_fn): {str(self.uncollate_fn)}\\n\"\n            f\"\\t(per_sample_transform): {str(self.per_sample_transform)}\\n\"\n            f\"\\t(serializer): {str(self.serializer)}\"\n        )\n\n\ndef default_uncollate(batch: Any):\n    \"\"\"\n    This function is used to uncollate a batch into samples.\n\n    Examples:\n        >>> a, b = default_uncollate(torch.rand((2,1)))\n    \"\"\"\n\n    batch_type = type(batch)\n\n    if isinstance(batch, Tensor):\n        if len(batch.shape) == 0:  # 0 shape tensors\n            return batch\n        return list(torch.unbind(batch, 0))\n\n    elif isinstance(batch, Mapping):\n        return [batch_type(dict(zip(batch, default_uncollate(t)))) for t in zip(*batch.values())]\n\n    elif isinstance(batch, tuple) and hasattr(batch, '_fields'):  # namedtuple\n        return [batch_type(*default_uncollate(sample)) for sample in zip(*batch)]\n\n    elif isinstance(batch, Sequence) and not isinstance(batch, str):\n        return [default_uncollate(sample) for sample in batch]\n\n    return batch\n", "path": "flash/core/data/batch.py"}]}