We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Serve sanity check never called with PL master
## 🐛 Bug

`LightningModule.run_sanity_check` has been renamed to `_run_sanity_check`

--- 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/trainer.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 import inspect
15 import warnings
16 from argparse import ArgumentParser, Namespace
17 from functools import wraps
18 from typing import Callable, List, Optional, Union
19 
20 import torch
21 from pytorch_lightning import LightningDataModule, LightningModule
22 from pytorch_lightning import Trainer as PlTrainer
23 from pytorch_lightning.callbacks import BaseFinetuning
24 from pytorch_lightning.utilities import rank_zero_warn
25 from pytorch_lightning.utilities.argparse import add_argparse_args, get_init_arguments_and_types, parse_env_variables
26 from pytorch_lightning.utilities.exceptions import MisconfigurationException
27 from torch.utils.data import DataLoader
28 
29 import flash
30 from flash.core.finetuning import _DEFAULTS_FINETUNE_STRATEGIES, instantiate_default_finetuning_callbacks
31 from flash.core.utilities.imports import _SERVE_AVAILABLE
32 
33 
34 def from_argparse_args(cls, args: Union[Namespace, ArgumentParser], **kwargs):
35     """Modified version of ``pytorch_lightning.utilities.argparse.from_argparse_args`` which populates ``valid_kwargs``
36     from ``pytorch_lightning.Trainer``."""
37     if isinstance(args, ArgumentParser):
38         args = cls.parse_argparser(args)
39 
40     params = vars(args)
41 
42     # we only want to pass in valid PLTrainer args, the rest may be user specific
43     valid_kwargs = inspect.signature(PlTrainer.__init__).parameters
44     trainer_kwargs = {name: params[name] for name in valid_kwargs if name in params}
45     trainer_kwargs.update(**kwargs)
46 
47     return cls(**trainer_kwargs)
48 
49 
50 def _defaults_from_env_vars(fn: Callable) -> Callable:
51     """Copy of ``pytorch_lightning.trainer.connectors.env_vars_connector._defaults_from_env_vars``. Required to fix
52     build error in readthedocs."""
53 
54     @wraps(fn)
55     def insert_env_defaults(self, *args, **kwargs):
56         cls = self.__class__  # get the class
57         if args:  # inace any args passed move them to kwargs
58             # parse only the argument names
59             cls_arg_names = [arg[0] for arg in get_init_arguments_and_types(cls)]
60             # convert args to kwargs
61             kwargs.update(dict(zip(cls_arg_names, args)))
62         env_variables = vars(parse_env_variables(cls))
63         # update the kwargs by env variables
64         kwargs = dict(list(env_variables.items()) + list(kwargs.items()))
65 
66         # all args were already moved to kwargs
67         return fn(self, **kwargs)
68 
69     return insert_env_defaults
70 
71 
72 class Trainer(PlTrainer):
73 
74     @_defaults_from_env_vars
75     def __init__(self, *args, serve_sanity_check: bool = False, **kwargs):
76         if flash._IS_TESTING:
77             if torch.cuda.is_available():
78                 kwargs["gpus"] = 1
79                 kwargs["max_epochs"] = 3
80                 kwargs["limit_train_batches"] = 1.0
81                 kwargs["limit_val_batches"] = 1.0
82                 kwargs["limit_test_batches"] = 1.0
83                 kwargs["fast_dev_run"] = False
84             else:
85                 kwargs["fast_dev_run"] = True
86         super().__init__(*args, **kwargs)
87 
88         self.serve_sanity_check = serve_sanity_check
89 
90     def run_sanity_check(self, ref_model):
91         super().run_sanity_check(ref_model)
92 
93         if self.serve_sanity_check and ref_model.is_servable and _SERVE_AVAILABLE:
94             ref_model.run_serve_sanity_check()
95 
96     def fit(
97         self,
98         model: LightningModule,
99         train_dataloader: Optional[DataLoader] = None,
100         val_dataloaders: Optional[Union[DataLoader, List[DataLoader]]] = None,
101         datamodule: Optional[LightningDataModule] = None,
102     ):
103         r"""
104         Runs the full optimization routine. Same as :meth:`pytorch_lightning.Trainer.fit`
105 
106         Args:
107             datamodule: A instance of :class:`LightningDataModule`.
108 
109             model: Model to fit.
110 
111             train_dataloader: A Pytorch DataLoader with training samples. If the model has
112                 a predefined train_dataloader method this will be skipped.
113 
114             val_dataloaders: Either a single Pytorch Dataloader or a list of them, specifying validation samples.
115                 If the model has a predefined val_dataloaders method this will be skipped
116         """
117         if any(isinstance(c, BaseFinetuning) for c in self.callbacks):
118             # TODO: if we find a finetuning callback in the trainer should we remove it? or just warn the user?
119             warnings.warn("Warning: You are calling fit(), but your trainer is using a fine-tuning callback")
120         return super().fit(model, train_dataloader, val_dataloaders, datamodule)
121 
122     def finetune(
123         self,
124         model: LightningModule,
125         train_dataloader: Optional[DataLoader] = None,
126         val_dataloaders: Optional[Union[DataLoader, List[DataLoader]]] = None,
127         datamodule: Optional[LightningDataModule] = None,
128         strategy: Optional[Union[str, BaseFinetuning]] = None,
129     ):
130         r"""
131 
132         Runs the full optimization routine. Same as :meth:`pytorch_lightning.Trainer.fit`, but unfreezes layers
133         of the backbone throughout training layers of the backbone throughout training.
134 
135         Args:
136             datamodule: A instance of :class:`LightningDataModule`.
137 
138             model: Model to fit.
139 
140             train_dataloader: A PyTorch DataLoader with training samples. If the model has
141                 a predefined train_dataloader method this will be skipped.
142 
143             val_dataloaders: Either a single PyTorch Dataloader or a list of them, specifying validation samples.
144                 If the model has a predefined val_dataloaders method this will be skipped
145 
146             strategy: Should either be a string or a finetuning callback subclassing
147                 :class:`pytorch_lightning.callbacks.BaseFinetuning`.
148 
149                 Default strategies can be enabled with these strings:
150 
151                 - ``"no_freeze"``,
152                 - ``"freeze"``,
153                 - ``"freeze_unfreeze"``,
154                 - ``"unfreeze_milestones"``.
155         """
156         self._resolve_callbacks(model, strategy)
157         return super().fit(model, train_dataloader, val_dataloaders, datamodule)
158 
159     def _resolve_callbacks(self, model, strategy):
160         """
161         This function is used to select the `BaseFinetuning` to be used for finetuning.
162         """
163         if strategy is not None and not isinstance(strategy, (str, BaseFinetuning)):
164             raise MisconfigurationException(
165                 "strategy should be a ``pytorch_lightning.callbacks.BaseFinetuning``"
166                 f"callback or a str within {list(_DEFAULTS_FINETUNE_STRATEGIES.keys())}"
167             )
168 
169         if isinstance(strategy, BaseFinetuning):
170             callback = [strategy]
171         else:
172             # todo: change to ``configure_callbacks`` when merged to Lightning.
173             model_callback = model.configure_finetune_callback()
174             if len(model_callback) > 1:
175                 raise MisconfigurationException(
176                     f"{model} configure_finetune_callback should create a list with only 1 callback"
177                 )
178             if len(model_callback) == 1:
179                 if strategy is not None:
180                     rank_zero_warn(
181                         "The model contains a default finetune callback. The provided {strategy} will be overriden.\n"
182                         " HINT: Provide a `BaseFinetuning` callback as strategy to make it prioritized. ", UserWarning
183                     )
184                 callback = model_callback
185             else:
186                 callback = instantiate_default_finetuning_callbacks(strategy)
187 
188         self.callbacks = self._merge_callbacks(self.callbacks, callback)
189 
190     @staticmethod
191     def _merge_callbacks(old_callbacks: List, new_callbacks: List) -> List:
192         """
193         This function keeps only 1 instance of each callback type,
194         extending new_callbacks with old_callbacks
195         """
196         if len(new_callbacks) == 0:
197             return old_callbacks
198         new_callbacks_types = {type(c) for c in new_callbacks}
199         old_callbacks_types = {type(c) for c in old_callbacks}
200         override_types = new_callbacks_types.intersection(old_callbacks_types)
201         new_callbacks.extend(c for c in old_callbacks if type(c) not in override_types)
202         return new_callbacks
203 
204     @classmethod
205     def add_argparse_args(cls, *args, **kwargs) -> ArgumentParser:
206         # the lightning trainer implementation does not support subclasses.
207         # context: https://github.com/PyTorchLightning/lightning-flash/issues/342#issuecomment-848892447
208         return add_argparse_args(PlTrainer, *args, **kwargs)
209 
210     @classmethod
211     def from_argparse_args(cls, args: Union[Namespace, ArgumentParser], **kwargs) -> 'Trainer':
212         # the lightning trainer implementation does not support subclasses.
213         # context: https://github.com/PyTorchLightning/lightning-flash/issues/342#issuecomment-848892447
214         return from_argparse_args(Trainer, args, **kwargs)
215
```
--- 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/trainer.py b/flash/core/trainer.py
--- a/flash/core/trainer.py
+++ b/flash/core/trainer.py
@@ -87,8 +87,15 @@
 
         self.serve_sanity_check = serve_sanity_check
 
+    def _run_sanity_check(self, ref_model):
+        if hasattr(super(), "_run_sanity_check"):
+            super()._run_sanity_check(ref_model)
+
+        self.run_sanity_check(ref_model)
+
     def run_sanity_check(self, ref_model):
-        super().run_sanity_check(ref_model)
+        if hasattr(super(), "run_sanity_check"):
+            super().run_sanity_check(ref_model)
 
         if self.serve_sanity_check and ref_model.is_servable and _SERVE_AVAILABLE:
             ref_model.run_serve_sanity_check()

{"golden_diff": "diff --git a/flash/core/trainer.py b/flash/core/trainer.py\n--- a/flash/core/trainer.py\n+++ b/flash/core/trainer.py\n@@ -87,8 +87,15 @@\n \n         self.serve_sanity_check = serve_sanity_check\n \n+    def _run_sanity_check(self, ref_model):\n+        if hasattr(super(), \"_run_sanity_check\"):\n+            super()._run_sanity_check(ref_model)\n+\n+        self.run_sanity_check(ref_model)\n+\n     def run_sanity_check(self, ref_model):\n-        super().run_sanity_check(ref_model)\n+        if hasattr(super(), \"run_sanity_check\"):\n+            super().run_sanity_check(ref_model)\n \n         if self.serve_sanity_check and ref_model.is_servable and _SERVE_AVAILABLE:\n             ref_model.run_serve_sanity_check()\n", "issue": "Serve sanity check never called with PL master\n## \ud83d\udc1b Bug\r\n\r\n`LightningModule.run_sanity_check` has been renamed to `_run_sanity_check`\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.\nimport inspect\nimport warnings\nfrom argparse import ArgumentParser, Namespace\nfrom functools import wraps\nfrom typing import Callable, List, Optional, Union\n\nimport torch\nfrom pytorch_lightning import LightningDataModule, LightningModule\nfrom pytorch_lightning import Trainer as PlTrainer\nfrom pytorch_lightning.callbacks import BaseFinetuning\nfrom pytorch_lightning.utilities import rank_zero_warn\nfrom pytorch_lightning.utilities.argparse import add_argparse_args, get_init_arguments_and_types, parse_env_variables\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\nfrom torch.utils.data import DataLoader\n\nimport flash\nfrom flash.core.finetuning import _DEFAULTS_FINETUNE_STRATEGIES, instantiate_default_finetuning_callbacks\nfrom flash.core.utilities.imports import _SERVE_AVAILABLE\n\n\ndef from_argparse_args(cls, args: Union[Namespace, ArgumentParser], **kwargs):\n    \"\"\"Modified version of ``pytorch_lightning.utilities.argparse.from_argparse_args`` which populates ``valid_kwargs``\n    from ``pytorch_lightning.Trainer``.\"\"\"\n    if isinstance(args, ArgumentParser):\n        args = cls.parse_argparser(args)\n\n    params = vars(args)\n\n    # we only want to pass in valid PLTrainer args, the rest may be user specific\n    valid_kwargs = inspect.signature(PlTrainer.__init__).parameters\n    trainer_kwargs = {name: params[name] for name in valid_kwargs if name in params}\n    trainer_kwargs.update(**kwargs)\n\n    return cls(**trainer_kwargs)\n\n\ndef _defaults_from_env_vars(fn: Callable) -> Callable:\n    \"\"\"Copy of ``pytorch_lightning.trainer.connectors.env_vars_connector._defaults_from_env_vars``. Required to fix\n    build error in readthedocs.\"\"\"\n\n    @wraps(fn)\n    def insert_env_defaults(self, *args, **kwargs):\n        cls = self.__class__  # get the class\n        if args:  # inace any args passed move them to kwargs\n            # parse only the argument names\n            cls_arg_names = [arg[0] for arg in get_init_arguments_and_types(cls)]\n            # convert args to kwargs\n            kwargs.update(dict(zip(cls_arg_names, args)))\n        env_variables = vars(parse_env_variables(cls))\n        # update the kwargs by env variables\n        kwargs = dict(list(env_variables.items()) + list(kwargs.items()))\n\n        # all args were already moved to kwargs\n        return fn(self, **kwargs)\n\n    return insert_env_defaults\n\n\nclass Trainer(PlTrainer):\n\n    @_defaults_from_env_vars\n    def __init__(self, *args, serve_sanity_check: bool = False, **kwargs):\n        if flash._IS_TESTING:\n            if torch.cuda.is_available():\n                kwargs[\"gpus\"] = 1\n                kwargs[\"max_epochs\"] = 3\n                kwargs[\"limit_train_batches\"] = 1.0\n                kwargs[\"limit_val_batches\"] = 1.0\n                kwargs[\"limit_test_batches\"] = 1.0\n                kwargs[\"fast_dev_run\"] = False\n            else:\n                kwargs[\"fast_dev_run\"] = True\n        super().__init__(*args, **kwargs)\n\n        self.serve_sanity_check = serve_sanity_check\n\n    def run_sanity_check(self, ref_model):\n        super().run_sanity_check(ref_model)\n\n        if self.serve_sanity_check and ref_model.is_servable and _SERVE_AVAILABLE:\n            ref_model.run_serve_sanity_check()\n\n    def fit(\n        self,\n        model: LightningModule,\n        train_dataloader: Optional[DataLoader] = None,\n        val_dataloaders: Optional[Union[DataLoader, List[DataLoader]]] = None,\n        datamodule: Optional[LightningDataModule] = None,\n    ):\n        r\"\"\"\n        Runs the full optimization routine. Same as :meth:`pytorch_lightning.Trainer.fit`\n\n        Args:\n            datamodule: A instance of :class:`LightningDataModule`.\n\n            model: Model to fit.\n\n            train_dataloader: A Pytorch DataLoader with training samples. If the model has\n                a predefined train_dataloader method this will be skipped.\n\n            val_dataloaders: Either a single Pytorch Dataloader or a list of them, specifying validation samples.\n                If the model has a predefined val_dataloaders method this will be skipped\n        \"\"\"\n        if any(isinstance(c, BaseFinetuning) for c in self.callbacks):\n            # TODO: if we find a finetuning callback in the trainer should we remove it? or just warn the user?\n            warnings.warn(\"Warning: You are calling fit(), but your trainer is using a fine-tuning callback\")\n        return super().fit(model, train_dataloader, val_dataloaders, datamodule)\n\n    def finetune(\n        self,\n        model: LightningModule,\n        train_dataloader: Optional[DataLoader] = None,\n        val_dataloaders: Optional[Union[DataLoader, List[DataLoader]]] = None,\n        datamodule: Optional[LightningDataModule] = None,\n        strategy: Optional[Union[str, BaseFinetuning]] = None,\n    ):\n        r\"\"\"\n\n        Runs the full optimization routine. Same as :meth:`pytorch_lightning.Trainer.fit`, but unfreezes layers\n        of the backbone throughout training layers of the backbone throughout training.\n\n        Args:\n            datamodule: A instance of :class:`LightningDataModule`.\n\n            model: Model to fit.\n\n            train_dataloader: A PyTorch DataLoader with training samples. If the model has\n                a predefined train_dataloader method this will be skipped.\n\n            val_dataloaders: Either a single PyTorch Dataloader or a list of them, specifying validation samples.\n                If the model has a predefined val_dataloaders method this will be skipped\n\n            strategy: Should either be a string or a finetuning callback subclassing\n                :class:`pytorch_lightning.callbacks.BaseFinetuning`.\n\n                Default strategies can be enabled with these strings:\n\n                - ``\"no_freeze\"``,\n                - ``\"freeze\"``,\n                - ``\"freeze_unfreeze\"``,\n                - ``\"unfreeze_milestones\"``.\n        \"\"\"\n        self._resolve_callbacks(model, strategy)\n        return super().fit(model, train_dataloader, val_dataloaders, datamodule)\n\n    def _resolve_callbacks(self, model, strategy):\n        \"\"\"\n        This function is used to select the `BaseFinetuning` to be used for finetuning.\n        \"\"\"\n        if strategy is not None and not isinstance(strategy, (str, BaseFinetuning)):\n            raise MisconfigurationException(\n                \"strategy should be a ``pytorch_lightning.callbacks.BaseFinetuning``\"\n                f\"callback or a str within {list(_DEFAULTS_FINETUNE_STRATEGIES.keys())}\"\n            )\n\n        if isinstance(strategy, BaseFinetuning):\n            callback = [strategy]\n        else:\n            # todo: change to ``configure_callbacks`` when merged to Lightning.\n            model_callback = model.configure_finetune_callback()\n            if len(model_callback) > 1:\n                raise MisconfigurationException(\n                    f\"{model} configure_finetune_callback should create a list with only 1 callback\"\n                )\n            if len(model_callback) == 1:\n                if strategy is not None:\n                    rank_zero_warn(\n                        \"The model contains a default finetune callback. The provided {strategy} will be overriden.\\n\"\n                        \" HINT: Provide a `BaseFinetuning` callback as strategy to make it prioritized. \", UserWarning\n                    )\n                callback = model_callback\n            else:\n                callback = instantiate_default_finetuning_callbacks(strategy)\n\n        self.callbacks = self._merge_callbacks(self.callbacks, callback)\n\n    @staticmethod\n    def _merge_callbacks(old_callbacks: List, new_callbacks: List) -> List:\n        \"\"\"\n        This function keeps only 1 instance of each callback type,\n        extending new_callbacks with old_callbacks\n        \"\"\"\n        if len(new_callbacks) == 0:\n            return old_callbacks\n        new_callbacks_types = {type(c) for c in new_callbacks}\n        old_callbacks_types = {type(c) for c in old_callbacks}\n        override_types = new_callbacks_types.intersection(old_callbacks_types)\n        new_callbacks.extend(c for c in old_callbacks if type(c) not in override_types)\n        return new_callbacks\n\n    @classmethod\n    def add_argparse_args(cls, *args, **kwargs) -> ArgumentParser:\n        # the lightning trainer implementation does not support subclasses.\n        # context: https://github.com/PyTorchLightning/lightning-flash/issues/342#issuecomment-848892447\n        return add_argparse_args(PlTrainer, *args, **kwargs)\n\n    @classmethod\n    def from_argparse_args(cls, args: Union[Namespace, ArgumentParser], **kwargs) -> 'Trainer':\n        # the lightning trainer implementation does not support subclasses.\n        # context: https://github.com/PyTorchLightning/lightning-flash/issues/342#issuecomment-848892447\n        return from_argparse_args(Trainer, args, **kwargs)\n", "path": "flash/core/trainer.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.\nimport inspect\nimport warnings\nfrom argparse import ArgumentParser, Namespace\nfrom functools import wraps\nfrom typing import Callable, List, Optional, Union\n\nimport torch\nfrom pytorch_lightning import LightningDataModule, LightningModule\nfrom pytorch_lightning import Trainer as PlTrainer\nfrom pytorch_lightning.callbacks import BaseFinetuning\nfrom pytorch_lightning.utilities import rank_zero_warn\nfrom pytorch_lightning.utilities.argparse import add_argparse_args, get_init_arguments_and_types, parse_env_variables\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\nfrom torch.utils.data import DataLoader\n\nimport flash\nfrom flash.core.finetuning import _DEFAULTS_FINETUNE_STRATEGIES, instantiate_default_finetuning_callbacks\nfrom flash.core.utilities.imports import _SERVE_AVAILABLE\n\n\ndef from_argparse_args(cls, args: Union[Namespace, ArgumentParser], **kwargs):\n    \"\"\"Modified version of ``pytorch_lightning.utilities.argparse.from_argparse_args`` which populates ``valid_kwargs``\n    from ``pytorch_lightning.Trainer``.\"\"\"\n    if isinstance(args, ArgumentParser):\n        args = cls.parse_argparser(args)\n\n    params = vars(args)\n\n    # we only want to pass in valid PLTrainer args, the rest may be user specific\n    valid_kwargs = inspect.signature(PlTrainer.__init__).parameters\n    trainer_kwargs = {name: params[name] for name in valid_kwargs if name in params}\n    trainer_kwargs.update(**kwargs)\n\n    return cls(**trainer_kwargs)\n\n\ndef _defaults_from_env_vars(fn: Callable) -> Callable:\n    \"\"\"Copy of ``pytorch_lightning.trainer.connectors.env_vars_connector._defaults_from_env_vars``. Required to fix\n    build error in readthedocs.\"\"\"\n\n    @wraps(fn)\n    def insert_env_defaults(self, *args, **kwargs):\n        cls = self.__class__  # get the class\n        if args:  # inace any args passed move them to kwargs\n            # parse only the argument names\n            cls_arg_names = [arg[0] for arg in get_init_arguments_and_types(cls)]\n            # convert args to kwargs\n            kwargs.update(dict(zip(cls_arg_names, args)))\n        env_variables = vars(parse_env_variables(cls))\n        # update the kwargs by env variables\n        kwargs = dict(list(env_variables.items()) + list(kwargs.items()))\n\n        # all args were already moved to kwargs\n        return fn(self, **kwargs)\n\n    return insert_env_defaults\n\n\nclass Trainer(PlTrainer):\n\n    @_defaults_from_env_vars\n    def __init__(self, *args, serve_sanity_check: bool = False, **kwargs):\n        if flash._IS_TESTING:\n            if torch.cuda.is_available():\n                kwargs[\"gpus\"] = 1\n                kwargs[\"max_epochs\"] = 3\n                kwargs[\"limit_train_batches\"] = 1.0\n                kwargs[\"limit_val_batches\"] = 1.0\n                kwargs[\"limit_test_batches\"] = 1.0\n                kwargs[\"fast_dev_run\"] = False\n            else:\n                kwargs[\"fast_dev_run\"] = True\n        super().__init__(*args, **kwargs)\n\n        self.serve_sanity_check = serve_sanity_check\n\n    def _run_sanity_check(self, ref_model):\n        if hasattr(super(), \"_run_sanity_check\"):\n            super()._run_sanity_check(ref_model)\n\n        self.run_sanity_check(ref_model)\n\n    def run_sanity_check(self, ref_model):\n        if hasattr(super(), \"run_sanity_check\"):\n            super().run_sanity_check(ref_model)\n\n        if self.serve_sanity_check and ref_model.is_servable and _SERVE_AVAILABLE:\n            ref_model.run_serve_sanity_check()\n\n    def fit(\n        self,\n        model: LightningModule,\n        train_dataloader: Optional[DataLoader] = None,\n        val_dataloaders: Optional[Union[DataLoader, List[DataLoader]]] = None,\n        datamodule: Optional[LightningDataModule] = None,\n    ):\n        r\"\"\"\n        Runs the full optimization routine. Same as :meth:`pytorch_lightning.Trainer.fit`\n\n        Args:\n            datamodule: A instance of :class:`LightningDataModule`.\n\n            model: Model to fit.\n\n            train_dataloader: A Pytorch DataLoader with training samples. If the model has\n                a predefined train_dataloader method this will be skipped.\n\n            val_dataloaders: Either a single Pytorch Dataloader or a list of them, specifying validation samples.\n                If the model has a predefined val_dataloaders method this will be skipped\n        \"\"\"\n        if any(isinstance(c, BaseFinetuning) for c in self.callbacks):\n            # TODO: if we find a finetuning callback in the trainer should we remove it? or just warn the user?\n            warnings.warn(\"Warning: You are calling fit(), but your trainer is using a fine-tuning callback\")\n        return super().fit(model, train_dataloader, val_dataloaders, datamodule)\n\n    def finetune(\n        self,\n        model: LightningModule,\n        train_dataloader: Optional[DataLoader] = None,\n        val_dataloaders: Optional[Union[DataLoader, List[DataLoader]]] = None,\n        datamodule: Optional[LightningDataModule] = None,\n        strategy: Optional[Union[str, BaseFinetuning]] = None,\n    ):\n        r\"\"\"\n\n        Runs the full optimization routine. Same as :meth:`pytorch_lightning.Trainer.fit`, but unfreezes layers\n        of the backbone throughout training layers of the backbone throughout training.\n\n        Args:\n            datamodule: A instance of :class:`LightningDataModule`.\n\n            model: Model to fit.\n\n            train_dataloader: A PyTorch DataLoader with training samples. If the model has\n                a predefined train_dataloader method this will be skipped.\n\n            val_dataloaders: Either a single PyTorch Dataloader or a list of them, specifying validation samples.\n                If the model has a predefined val_dataloaders method this will be skipped\n\n            strategy: Should either be a string or a finetuning callback subclassing\n                :class:`pytorch_lightning.callbacks.BaseFinetuning`.\n\n                Default strategies can be enabled with these strings:\n\n                - ``\"no_freeze\"``,\n                - ``\"freeze\"``,\n                - ``\"freeze_unfreeze\"``,\n                - ``\"unfreeze_milestones\"``.\n        \"\"\"\n        self._resolve_callbacks(model, strategy)\n        return super().fit(model, train_dataloader, val_dataloaders, datamodule)\n\n    def _resolve_callbacks(self, model, strategy):\n        \"\"\"\n        This function is used to select the `BaseFinetuning` to be used for finetuning.\n        \"\"\"\n        if strategy is not None and not isinstance(strategy, (str, BaseFinetuning)):\n            raise MisconfigurationException(\n                \"strategy should be a ``pytorch_lightning.callbacks.BaseFinetuning``\"\n                f\"callback or a str within {list(_DEFAULTS_FINETUNE_STRATEGIES.keys())}\"\n            )\n\n        if isinstance(strategy, BaseFinetuning):\n            callback = [strategy]\n        else:\n            # todo: change to ``configure_callbacks`` when merged to Lightning.\n            model_callback = model.configure_finetune_callback()\n            if len(model_callback) > 1:\n                raise MisconfigurationException(\n                    f\"{model} configure_finetune_callback should create a list with only 1 callback\"\n                )\n            if len(model_callback) == 1:\n                if strategy is not None:\n                    rank_zero_warn(\n                        \"The model contains a default finetune callback. The provided {strategy} will be overriden.\\n\"\n                        \" HINT: Provide a `BaseFinetuning` callback as strategy to make it prioritized. \", UserWarning\n                    )\n                callback = model_callback\n            else:\n                callback = instantiate_default_finetuning_callbacks(strategy)\n\n        self.callbacks = self._merge_callbacks(self.callbacks, callback)\n\n    @staticmethod\n    def _merge_callbacks(old_callbacks: List, new_callbacks: List) -> List:\n        \"\"\"\n        This function keeps only 1 instance of each callback type,\n        extending new_callbacks with old_callbacks\n        \"\"\"\n        if len(new_callbacks) == 0:\n            return old_callbacks\n        new_callbacks_types = {type(c) for c in new_callbacks}\n        old_callbacks_types = {type(c) for c in old_callbacks}\n        override_types = new_callbacks_types.intersection(old_callbacks_types)\n        new_callbacks.extend(c for c in old_callbacks if type(c) not in override_types)\n        return new_callbacks\n\n    @classmethod\n    def add_argparse_args(cls, *args, **kwargs) -> ArgumentParser:\n        # the lightning trainer implementation does not support subclasses.\n        # context: https://github.com/PyTorchLightning/lightning-flash/issues/342#issuecomment-848892447\n        return add_argparse_args(PlTrainer, *args, **kwargs)\n\n    @classmethod\n    def from_argparse_args(cls, args: Union[Namespace, ArgumentParser], **kwargs) -> 'Trainer':\n        # the lightning trainer implementation does not support subclasses.\n        # context: https://github.com/PyTorchLightning/lightning-flash/issues/342#issuecomment-848892447\n        return from_argparse_args(Trainer, args, **kwargs)\n", "path": "flash/core/trainer.py"}]}