We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Breaking compatibility with custom datatypes implementing `.to`
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->
Bring back compatibility for custom datatypes in collections implementing `.to` for transferring data.

### Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
I am using Pytorch Lightning together with Pytorch Geometric. Pytorch Geometric implements several custom datatypes and dataloaders which is really useful for geometric deep learning. Everything worked well with pytorch lightning 0.7.6, as the custom datatypes implement a `.to` method for transferring the data to different devices.
However, with the recent 0.8.1 update, this is no longer possible and I had to scour the documentation to be able to implement a fix using `transfer_batch_to_device(batch, device)`. This is in my opinion not very pretty, as my batch looks like this 
```
{"data": pytorch geometric batch object, "id": tensor, ...}
```
i.e. it is just a dictionary of types that all implement the `.to` method.

### Pitch

<!-- A clear and concise description of what you want to happen. -->

- Make it possible for classes implementing the `.to` method to be transferred automatically
- If part of the batch could not be transferred automatically output a warning letting the user know, that a custom transfer function for the batch might be required, or to implement the `.to` method for custom datatypes in the batch
- Add a note to the introduction guide about custom datatypes and handling for custom datatypes

### Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->
If this change was intentional and the behavior of trying to call the `.to` method is not desired, I think there should definitely be some more documentation about this, in a more obvious place.

### Additional context

<!-- Add any other context or screenshots about the feature request here. -->


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pytorch_lightning/core/hooks.py`
Content:
```
1 from typing import Any
2 
3 import torch
4 from torch import Tensor
5 from torch.nn import Module
6 from torch.optim.optimizer import Optimizer
7 from pytorch_lightning.utilities import move_data_to_device
8 
9 
10 try:
11     from apex import amp
12 except ImportError:
13     APEX_AVAILABLE = False
14 else:
15     APEX_AVAILABLE = True
16 
17 
18 class ModelHooks(Module):
19 
20     def setup(self, stage: str):
21         """
22         Called at the beginning of fit and test.
23         This is a good hook when you need to build models dynamically or adjust something about them.
24         This hook is called on every process when using DDP.
25 
26         Args:
27             stage: either 'fit' or 'test'
28 
29         Example::
30 
31             class LitModel(...):
32                 def __init__(self):
33                     self.l1 = None
34 
35                 def prepare_data(self):
36                     download_data()
37                     tokenize()
38 
39                     # don't do this
40                     self.something = else
41 
42                 def setup(step):
43                     data = Load_data(...)
44                     self.l1 = nn.Linear(28, data.num_classes)
45 
46         """
47 
48     def teardown(self, stage: str):
49         """
50         Called at the end of fit and test.
51 
52         Args:
53             stage: either 'fit' or 'test'
54         """
55 
56     def on_fit_start(self):
57         """
58         Called at the very beginning of fit.
59         If on DDP it is called on every process
60         """
61 
62     def on_fit_end(self):
63         """
64         Called at the very end of fit.
65         If on DDP it is called on every process
66         """
67 
68     # TODO: remove in v0.9.0
69     def on_sanity_check_start(self):
70         """
71         Called before starting evaluation.
72 
73         Warning:
74             Deprecated. Will be removed in v0.9.0.
75         """
76 
77     def on_train_start(self) -> None:
78         """
79         Called at the beginning of training before sanity check.
80         """
81         # do something at the start of training
82 
83     def on_train_end(self) -> None:
84         """
85         Called at the end of training before logger experiment is closed.
86         """
87         # do something at the end of training
88 
89     def on_batch_start(self, batch: Any) -> None:
90         """
91         Called in the training loop before anything happens for that batch.
92 
93         If you return -1 here, you will skip training for the rest of the current epoch.
94 
95         Args:
96             batch: The batched data as it is returned by the training DataLoader.
97         """
98         # do something when the batch starts
99 
100     def on_batch_end(self) -> None:
101         """
102         Called in the training loop after the batch.
103         """
104         # do something when the batch ends
105 
106     def on_epoch_start(self) -> None:
107         """
108         Called in the training loop at the very beginning of the epoch.
109         """
110         # do something when the epoch starts
111 
112     def on_epoch_end(self) -> None:
113         """
114         Called in the training loop at the very end of the epoch.
115         """
116         # do something when the epoch ends
117 
118     def on_pre_performance_check(self) -> None:
119         """
120         Called at the very beginning of the validation loop.
121         """
122         # do something before validation starts
123 
124     def on_post_performance_check(self) -> None:
125         """
126         Called at the very end of the validation loop.
127         """
128         # do something before validation end
129 
130     def on_before_zero_grad(self, optimizer: Optimizer) -> None:
131         """
132         Called after optimizer.step() and before optimizer.zero_grad().
133 
134         Called in the training loop after taking an optimizer step and before zeroing grads.
135         Good place to inspect weight information with weights updated.
136 
137         This is where it is called::
138 
139             for optimizer in optimizers:
140                 optimizer.step()
141                 model.on_before_zero_grad(optimizer) # < ---- called here
142                 optimizer.zero_grad
143 
144         Args:
145             optimizer: The optimizer for which grads should be zeroed.
146         """
147         # do something with the optimizer or inspect it.
148 
149     def on_after_backward(self) -> None:
150         """
151         Called in the training loop after loss.backward() and before optimizers do anything.
152         This is the ideal place to inspect or log gradient information.
153 
154         Example::
155 
156             def on_after_backward(self):
157                 # example to inspect gradient information in tensorboard
158                 if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
159                     params = self.state_dict()
160                     for k, v in params.items():
161                         grads = v
162                         name = k
163                         self.logger.experiment.add_histogram(tag=name, values=grads,
164                                                              global_step=self.trainer.global_step)
165 
166         """
167 
168     def backward(self, trainer, loss: Tensor, optimizer: Optimizer, optimizer_idx: int) -> None:
169         """
170         Override backward with your own implementation if you need to.
171 
172         Args:
173             trainer: Pointer to the trainer
174             loss: Loss is already scaled by accumulated grads
175             optimizer: Current optimizer being used
176             optimizer_idx: Index of the current optimizer being used
177 
178         Called to perform backward step.
179         Feel free to override as needed.
180 
181         The loss passed in has already been scaled for accumulated gradients if requested.
182 
183         Example::
184 
185             def backward(self, trainer, loss, optimizer, optimizer_idx):
186                 loss.backward()
187 
188         """
189         loss.backward()
190 
191     def amp_scale_loss(self, unscaled_loss, optimizer, optimizer_idx):
192         if self.trainer.use_native_amp:
193             scaled_loss = self.trainer.scaler.scale(unscaled_loss)
194 
195         else:
196             scaled_loss = amp.scale_loss(unscaled_loss, optimizer)
197 
198         return scaled_loss
199 
200     def transfer_batch_to_device(self, batch: Any, device: torch.device) -> Any:
201         """
202         Override this hook if your :class:`~torch.utils.data.DataLoader` returns tensors
203         wrapped in a custom data structure.
204 
205         The data types listed below (and any arbitrary nesting of them) are supported out of the box:
206 
207         - :class:`torch.Tensor`
208         - :class:`list`
209         - :class:`dict`
210         - :class:`tuple`
211         - ``torchtext.data.Batch`` (COMING SOON)
212 
213         For anything else, you need to define how the data is moved to the target device (CPU, GPU, TPU, ...).
214 
215         Example::
216 
217             def transfer_batch_to_device(self, batch, device)
218                 if isinstance(batch, CustomBatch):
219                     # move all tensors in your custom data structure to the device
220                     batch.samples = batch.samples.to(device)
221                     batch.targets = batch.targets.to(device)
222                 else:
223                     batch = super().transfer_batch_to_device(data, device)
224                 return batch
225 
226         Args:
227             batch: A batch of data that needs to be transferred to a new device.
228             device: The target device as defined in PyTorch.
229 
230         Returns:
231             A reference to the data on the new device.
232 
233         Note:
234             This hook should only transfer the data and not modify it, nor should it move the data to
235             any other device than the one passed in as argument (unless you know what you are doing).
236             The :class:`~pytorch_lightning.trainer.trainer.Trainer` already takes care of splitting the
237             batch and determines the target devices.
238 
239         See Also:
240             - :func:`~pytorch_lightning.utilities.apply_func.move_data_to_device`
241             - :func:`~pytorch_lightning.utilities.apply_func.apply_to_collection`
242         """
243         return move_data_to_device(batch, device)
244
```

Path: `pytorch_lightning/utilities/apply_func.py`
Content:
```
1 from collections import Mapping, Sequence
2 from typing import Any, Callable, Union
3 
4 import torch
5 
6 
7 def apply_to_collection(data: Any, dtype: Union[type, tuple], function: Callable, *args, **kwargs) -> Any:
8     """
9     Recursively applies a function to all elements of a certain dtype.
10 
11     Args:
12         data: the collection to apply the function to
13         dtype: the given function will be applied to all elements of this dtype
14         function: the function to apply
15         *args: positional arguments (will be forwarded to calls of ``function``)
16         **kwargs: keyword arguments (will be forwarded to calls of ``function``)
17 
18     Returns:
19         the resulting collection
20 
21     """
22     elem_type = type(data)
23 
24     # Breaking condition
25     if isinstance(data, dtype):
26         return function(data, *args, **kwargs)
27 
28     # Recursively apply to collection items
29     elif isinstance(data, Mapping):
30         return elem_type({k: apply_to_collection(v, dtype, function, *args, **kwargs)
31                           for k, v in data.items()})
32     elif isinstance(data, tuple) and hasattr(data, '_fields'):  # named tuple
33         return elem_type(*(apply_to_collection(d, dtype, function, *args, **kwargs) for d in data))
34     elif isinstance(data, Sequence) and not isinstance(data, str):
35         return elem_type([apply_to_collection(d, dtype, function, *args, **kwargs) for d in data])
36 
37     # data is neither of dtype, nor a collection
38     return data
39 
40 
41 def move_data_to_device(batch: Any, device: torch.device):
42     """
43     Transfers a collection of tensors to the given device.
44 
45     Args:
46         batch: A tensor or collection of tensors. See :func:`apply_to_collection`
47             for a list of supported collection types.
48         device: The device to which tensors should be moved
49 
50     Return:
51         the same collection but with all contained tensors residing on the new device.
52 
53     See Also:
54         - :meth:`torch.Tensor.to`
55         - :class:`torch.device`
56     """
57     def to(tensor):
58         return tensor.to(device, non_blocking=True)
59     return apply_to_collection(batch, dtype=torch.Tensor, function=to)
60
```
--- 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/pytorch_lightning/core/hooks.py b/pytorch_lightning/core/hooks.py
--- a/pytorch_lightning/core/hooks.py
+++ b/pytorch_lightning/core/hooks.py
@@ -204,7 +204,7 @@
 
         The data types listed below (and any arbitrary nesting of them) are supported out of the box:
 
-        - :class:`torch.Tensor`
+        - :class:`torch.Tensor` or anything that implements `.to(...)`
         - :class:`list`
         - :class:`dict`
         - :class:`tuple`
diff --git a/pytorch_lightning/utilities/apply_func.py b/pytorch_lightning/utilities/apply_func.py
--- a/pytorch_lightning/utilities/apply_func.py
+++ b/pytorch_lightning/utilities/apply_func.py
@@ -1,3 +1,4 @@
+from abc import ABC
 from collections import Mapping, Sequence
 from typing import Any, Callable, Union
 
@@ -38,14 +39,43 @@
     return data
 
 
+class TransferableDataType(ABC):
+    """
+    A custom type for data that can be moved to a torch device via `.to(...)`.
+
+    Example:
+
+        >>> isinstance(dict, TransferableDataType)
+        False
+        >>> isinstance(torch.rand(2, 3), TransferableDataType)
+        True
+        >>> class CustomObject:
+        ...     def __init__(self):
+        ...         self.x = torch.rand(2, 2)
+        ...     def to(self, device):
+        ...         self.x = self.x.to(device)
+        ...         return self
+        >>> isinstance(CustomObject(), TransferableDataType)
+        True
+    """
+
+    @classmethod
+    def __subclasshook__(cls, subclass):
+        if cls is TransferableDataType:
+            to = getattr(subclass, "to", None)
+            return callable(to)
+        return NotImplemented
+
+
 def move_data_to_device(batch: Any, device: torch.device):
     """
-    Transfers a collection of tensors to the given device.
+    Transfers a collection of data to the given device. Any object that defines a method
+    ``to(device)`` will be moved and all other objects in the collection will be left untouched.
 
     Args:
-        batch: A tensor or collection of tensors. See :func:`apply_to_collection`
-            for a list of supported collection types.
-        device: The device to which tensors should be moved
+        batch: A tensor or collection of tensors or anything that has a method `.to(...)`.
+            See :func:`apply_to_collection` for a list of supported collection types.
+        device: The device to which the data should be moved
 
     Return:
         the same collection but with all contained tensors residing on the new device.
@@ -54,6 +84,6 @@
         - :meth:`torch.Tensor.to`
         - :class:`torch.device`
     """
-    def to(tensor):
-        return tensor.to(device, non_blocking=True)
-    return apply_to_collection(batch, dtype=torch.Tensor, function=to)
+    def to(data):
+        return data.to(device, non_blocking=True)
+    return apply_to_collection(batch, dtype=TransferableDataType, function=to)

{"golden_diff": "diff --git a/pytorch_lightning/core/hooks.py b/pytorch_lightning/core/hooks.py\n--- a/pytorch_lightning/core/hooks.py\n+++ b/pytorch_lightning/core/hooks.py\n@@ -204,7 +204,7 @@\n \n         The data types listed below (and any arbitrary nesting of them) are supported out of the box:\n \n-        - :class:`torch.Tensor`\n+        - :class:`torch.Tensor` or anything that implements `.to(...)`\n         - :class:`list`\n         - :class:`dict`\n         - :class:`tuple`\ndiff --git a/pytorch_lightning/utilities/apply_func.py b/pytorch_lightning/utilities/apply_func.py\n--- a/pytorch_lightning/utilities/apply_func.py\n+++ b/pytorch_lightning/utilities/apply_func.py\n@@ -1,3 +1,4 @@\n+from abc import ABC\n from collections import Mapping, Sequence\n from typing import Any, Callable, Union\n \n@@ -38,14 +39,43 @@\n     return data\n \n \n+class TransferableDataType(ABC):\n+    \"\"\"\n+    A custom type for data that can be moved to a torch device via `.to(...)`.\n+\n+    Example:\n+\n+        >>> isinstance(dict, TransferableDataType)\n+        False\n+        >>> isinstance(torch.rand(2, 3), TransferableDataType)\n+        True\n+        >>> class CustomObject:\n+        ...     def __init__(self):\n+        ...         self.x = torch.rand(2, 2)\n+        ...     def to(self, device):\n+        ...         self.x = self.x.to(device)\n+        ...         return self\n+        >>> isinstance(CustomObject(), TransferableDataType)\n+        True\n+    \"\"\"\n+\n+    @classmethod\n+    def __subclasshook__(cls, subclass):\n+        if cls is TransferableDataType:\n+            to = getattr(subclass, \"to\", None)\n+            return callable(to)\n+        return NotImplemented\n+\n+\n def move_data_to_device(batch: Any, device: torch.device):\n     \"\"\"\n-    Transfers a collection of tensors to the given device.\n+    Transfers a collection of data to the given device. Any object that defines a method\n+    ``to(device)`` will be moved and all other objects in the collection will be left untouched.\n \n     Args:\n-        batch: A tensor or collection of tensors. See :func:`apply_to_collection`\n-            for a list of supported collection types.\n-        device: The device to which tensors should be moved\n+        batch: A tensor or collection of tensors or anything that has a method `.to(...)`.\n+            See :func:`apply_to_collection` for a list of supported collection types.\n+        device: The device to which the data should be moved\n \n     Return:\n         the same collection but with all contained tensors residing on the new device.\n@@ -54,6 +84,6 @@\n         - :meth:`torch.Tensor.to`\n         - :class:`torch.device`\n     \"\"\"\n-    def to(tensor):\n-        return tensor.to(device, non_blocking=True)\n-    return apply_to_collection(batch, dtype=torch.Tensor, function=to)\n+    def to(data):\n+        return data.to(device, non_blocking=True)\n+    return apply_to_collection(batch, dtype=TransferableDataType, function=to)\n", "issue": "Breaking compatibility with custom datatypes implementing `.to`\n## \ud83d\ude80 Feature\r\n<!-- A clear and concise description of the feature proposal -->\r\nBring back compatibility for custom datatypes in collections implementing `.to` for transferring data.\r\n\r\n### Motivation\r\n\r\n<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->\r\nI am using Pytorch Lightning together with Pytorch Geometric. Pytorch Geometric implements several custom datatypes and dataloaders which is really useful for geometric deep learning. Everything worked well with pytorch lightning 0.7.6, as the custom datatypes implement a `.to` method for transferring the data to different devices.\r\nHowever, with the recent 0.8.1 update, this is no longer possible and I had to scour the documentation to be able to implement a fix using `transfer_batch_to_device(batch, device)`. This is in my opinion not very pretty, as my batch looks like this \r\n```\r\n{\"data\": pytorch geometric batch object, \"id\": tensor, ...}\r\n```\r\ni.e. it is just a dictionary of types that all implement the `.to` method.\r\n\r\n### Pitch\r\n\r\n<!-- A clear and concise description of what you want to happen. -->\r\n\r\n- Make it possible for classes implementing the `.to` method to be transferred automatically\r\n- If part of the batch could not be transferred automatically output a warning letting the user know, that a custom transfer function for the batch might be required, or to implement the `.to` method for custom datatypes in the batch\r\n- Add a note to the introduction guide about custom datatypes and handling for custom datatypes\r\n\r\n### Alternatives\r\n\r\n<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->\r\nIf this change was intentional and the behavior of trying to call the `.to` method is not desired, I think there should definitely be some more documentation about this, in a more obvious place.\r\n\r\n### Additional context\r\n\r\n<!-- Add any other context or screenshots about the feature request here. -->\r\n\n", "before_files": [{"content": "from typing import Any\n\nimport torch\nfrom torch import Tensor\nfrom torch.nn import Module\nfrom torch.optim.optimizer import Optimizer\nfrom pytorch_lightning.utilities import move_data_to_device\n\n\ntry:\n    from apex import amp\nexcept ImportError:\n    APEX_AVAILABLE = False\nelse:\n    APEX_AVAILABLE = True\n\n\nclass ModelHooks(Module):\n\n    def setup(self, stage: str):\n        \"\"\"\n        Called at the beginning of fit and test.\n        This is a good hook when you need to build models dynamically or adjust something about them.\n        This hook is called on every process when using DDP.\n\n        Args:\n            stage: either 'fit' or 'test'\n\n        Example::\n\n            class LitModel(...):\n                def __init__(self):\n                    self.l1 = None\n\n                def prepare_data(self):\n                    download_data()\n                    tokenize()\n\n                    # don't do this\n                    self.something = else\n\n                def setup(step):\n                    data = Load_data(...)\n                    self.l1 = nn.Linear(28, data.num_classes)\n\n        \"\"\"\n\n    def teardown(self, stage: str):\n        \"\"\"\n        Called at the end of fit and test.\n\n        Args:\n            stage: either 'fit' or 'test'\n        \"\"\"\n\n    def on_fit_start(self):\n        \"\"\"\n        Called at the very beginning of fit.\n        If on DDP it is called on every process\n        \"\"\"\n\n    def on_fit_end(self):\n        \"\"\"\n        Called at the very end of fit.\n        If on DDP it is called on every process\n        \"\"\"\n\n    # TODO: remove in v0.9.0\n    def on_sanity_check_start(self):\n        \"\"\"\n        Called before starting evaluation.\n\n        Warning:\n            Deprecated. Will be removed in v0.9.0.\n        \"\"\"\n\n    def on_train_start(self) -> None:\n        \"\"\"\n        Called at the beginning of training before sanity check.\n        \"\"\"\n        # do something at the start of training\n\n    def on_train_end(self) -> None:\n        \"\"\"\n        Called at the end of training before logger experiment is closed.\n        \"\"\"\n        # do something at the end of training\n\n    def on_batch_start(self, batch: Any) -> None:\n        \"\"\"\n        Called in the training loop before anything happens for that batch.\n\n        If you return -1 here, you will skip training for the rest of the current epoch.\n\n        Args:\n            batch: The batched data as it is returned by the training DataLoader.\n        \"\"\"\n        # do something when the batch starts\n\n    def on_batch_end(self) -> None:\n        \"\"\"\n        Called in the training loop after the batch.\n        \"\"\"\n        # do something when the batch ends\n\n    def on_epoch_start(self) -> None:\n        \"\"\"\n        Called in the training loop at the very beginning of the epoch.\n        \"\"\"\n        # do something when the epoch starts\n\n    def on_epoch_end(self) -> None:\n        \"\"\"\n        Called in the training loop at the very end of the epoch.\n        \"\"\"\n        # do something when the epoch ends\n\n    def on_pre_performance_check(self) -> None:\n        \"\"\"\n        Called at the very beginning of the validation loop.\n        \"\"\"\n        # do something before validation starts\n\n    def on_post_performance_check(self) -> None:\n        \"\"\"\n        Called at the very end of the validation loop.\n        \"\"\"\n        # do something before validation end\n\n    def on_before_zero_grad(self, optimizer: Optimizer) -> None:\n        \"\"\"\n        Called after optimizer.step() and before optimizer.zero_grad().\n\n        Called in the training loop after taking an optimizer step and before zeroing grads.\n        Good place to inspect weight information with weights updated.\n\n        This is where it is called::\n\n            for optimizer in optimizers:\n                optimizer.step()\n                model.on_before_zero_grad(optimizer) # < ---- called here\n                optimizer.zero_grad\n\n        Args:\n            optimizer: The optimizer for which grads should be zeroed.\n        \"\"\"\n        # do something with the optimizer or inspect it.\n\n    def on_after_backward(self) -> None:\n        \"\"\"\n        Called in the training loop after loss.backward() and before optimizers do anything.\n        This is the ideal place to inspect or log gradient information.\n\n        Example::\n\n            def on_after_backward(self):\n                # example to inspect gradient information in tensorboard\n                if self.trainer.global_step % 25 == 0:  # don't make the tf file huge\n                    params = self.state_dict()\n                    for k, v in params.items():\n                        grads = v\n                        name = k\n                        self.logger.experiment.add_histogram(tag=name, values=grads,\n                                                             global_step=self.trainer.global_step)\n\n        \"\"\"\n\n    def backward(self, trainer, loss: Tensor, optimizer: Optimizer, optimizer_idx: int) -> None:\n        \"\"\"\n        Override backward with your own implementation if you need to.\n\n        Args:\n            trainer: Pointer to the trainer\n            loss: Loss is already scaled by accumulated grads\n            optimizer: Current optimizer being used\n            optimizer_idx: Index of the current optimizer being used\n\n        Called to perform backward step.\n        Feel free to override as needed.\n\n        The loss passed in has already been scaled for accumulated gradients if requested.\n\n        Example::\n\n            def backward(self, trainer, loss, optimizer, optimizer_idx):\n                loss.backward()\n\n        \"\"\"\n        loss.backward()\n\n    def amp_scale_loss(self, unscaled_loss, optimizer, optimizer_idx):\n        if self.trainer.use_native_amp:\n            scaled_loss = self.trainer.scaler.scale(unscaled_loss)\n\n        else:\n            scaled_loss = amp.scale_loss(unscaled_loss, optimizer)\n\n        return scaled_loss\n\n    def transfer_batch_to_device(self, batch: Any, device: torch.device) -> Any:\n        \"\"\"\n        Override this hook if your :class:`~torch.utils.data.DataLoader` returns tensors\n        wrapped in a custom data structure.\n\n        The data types listed below (and any arbitrary nesting of them) are supported out of the box:\n\n        - :class:`torch.Tensor`\n        - :class:`list`\n        - :class:`dict`\n        - :class:`tuple`\n        - ``torchtext.data.Batch`` (COMING SOON)\n\n        For anything else, you need to define how the data is moved to the target device (CPU, GPU, TPU, ...).\n\n        Example::\n\n            def transfer_batch_to_device(self, batch, device)\n                if isinstance(batch, CustomBatch):\n                    # move all tensors in your custom data structure to the device\n                    batch.samples = batch.samples.to(device)\n                    batch.targets = batch.targets.to(device)\n                else:\n                    batch = super().transfer_batch_to_device(data, device)\n                return batch\n\n        Args:\n            batch: A batch of data that needs to be transferred to a new device.\n            device: The target device as defined in PyTorch.\n\n        Returns:\n            A reference to the data on the new device.\n\n        Note:\n            This hook should only transfer the data and not modify it, nor should it move the data to\n            any other device than the one passed in as argument (unless you know what you are doing).\n            The :class:`~pytorch_lightning.trainer.trainer.Trainer` already takes care of splitting the\n            batch and determines the target devices.\n\n        See Also:\n            - :func:`~pytorch_lightning.utilities.apply_func.move_data_to_device`\n            - :func:`~pytorch_lightning.utilities.apply_func.apply_to_collection`\n        \"\"\"\n        return move_data_to_device(batch, device)\n", "path": "pytorch_lightning/core/hooks.py"}, {"content": "from collections import Mapping, Sequence\nfrom typing import Any, Callable, Union\n\nimport torch\n\n\ndef apply_to_collection(data: Any, dtype: Union[type, tuple], function: Callable, *args, **kwargs) -> Any:\n    \"\"\"\n    Recursively applies a function to all elements of a certain dtype.\n\n    Args:\n        data: the collection to apply the function to\n        dtype: the given function will be applied to all elements of this dtype\n        function: the function to apply\n        *args: positional arguments (will be forwarded to calls of ``function``)\n        **kwargs: keyword arguments (will be forwarded to calls of ``function``)\n\n    Returns:\n        the resulting collection\n\n    \"\"\"\n    elem_type = type(data)\n\n    # Breaking condition\n    if isinstance(data, dtype):\n        return function(data, *args, **kwargs)\n\n    # Recursively apply to collection items\n    elif isinstance(data, Mapping):\n        return elem_type({k: apply_to_collection(v, dtype, function, *args, **kwargs)\n                          for k, v in data.items()})\n    elif isinstance(data, tuple) and hasattr(data, '_fields'):  # named tuple\n        return elem_type(*(apply_to_collection(d, dtype, function, *args, **kwargs) for d in data))\n    elif isinstance(data, Sequence) and not isinstance(data, str):\n        return elem_type([apply_to_collection(d, dtype, function, *args, **kwargs) for d in data])\n\n    # data is neither of dtype, nor a collection\n    return data\n\n\ndef move_data_to_device(batch: Any, device: torch.device):\n    \"\"\"\n    Transfers a collection of tensors to the given device.\n\n    Args:\n        batch: A tensor or collection of tensors. See :func:`apply_to_collection`\n            for a list of supported collection types.\n        device: The device to which tensors should be moved\n\n    Return:\n        the same collection but with all contained tensors residing on the new device.\n\n    See Also:\n        - :meth:`torch.Tensor.to`\n        - :class:`torch.device`\n    \"\"\"\n    def to(tensor):\n        return tensor.to(device, non_blocking=True)\n    return apply_to_collection(batch, dtype=torch.Tensor, function=to)\n", "path": "pytorch_lightning/utilities/apply_func.py"}], "after_files": [{"content": "from typing import Any\n\nimport torch\nfrom torch import Tensor\nfrom torch.nn import Module\nfrom torch.optim.optimizer import Optimizer\nfrom pytorch_lightning.utilities import move_data_to_device\n\n\ntry:\n    from apex import amp\nexcept ImportError:\n    APEX_AVAILABLE = False\nelse:\n    APEX_AVAILABLE = True\n\n\nclass ModelHooks(Module):\n\n    def setup(self, stage: str):\n        \"\"\"\n        Called at the beginning of fit and test.\n        This is a good hook when you need to build models dynamically or adjust something about them.\n        This hook is called on every process when using DDP.\n\n        Args:\n            stage: either 'fit' or 'test'\n\n        Example::\n\n            class LitModel(...):\n                def __init__(self):\n                    self.l1 = None\n\n                def prepare_data(self):\n                    download_data()\n                    tokenize()\n\n                    # don't do this\n                    self.something = else\n\n                def setup(step):\n                    data = Load_data(...)\n                    self.l1 = nn.Linear(28, data.num_classes)\n\n        \"\"\"\n\n    def teardown(self, stage: str):\n        \"\"\"\n        Called at the end of fit and test.\n\n        Args:\n            stage: either 'fit' or 'test'\n        \"\"\"\n\n    def on_fit_start(self):\n        \"\"\"\n        Called at the very beginning of fit.\n        If on DDP it is called on every process\n        \"\"\"\n\n    def on_fit_end(self):\n        \"\"\"\n        Called at the very end of fit.\n        If on DDP it is called on every process\n        \"\"\"\n\n    # TODO: remove in v0.9.0\n    def on_sanity_check_start(self):\n        \"\"\"\n        Called before starting evaluation.\n\n        Warning:\n            Deprecated. Will be removed in v0.9.0.\n        \"\"\"\n\n    def on_train_start(self) -> None:\n        \"\"\"\n        Called at the beginning of training before sanity check.\n        \"\"\"\n        # do something at the start of training\n\n    def on_train_end(self) -> None:\n        \"\"\"\n        Called at the end of training before logger experiment is closed.\n        \"\"\"\n        # do something at the end of training\n\n    def on_batch_start(self, batch: Any) -> None:\n        \"\"\"\n        Called in the training loop before anything happens for that batch.\n\n        If you return -1 here, you will skip training for the rest of the current epoch.\n\n        Args:\n            batch: The batched data as it is returned by the training DataLoader.\n        \"\"\"\n        # do something when the batch starts\n\n    def on_batch_end(self) -> None:\n        \"\"\"\n        Called in the training loop after the batch.\n        \"\"\"\n        # do something when the batch ends\n\n    def on_epoch_start(self) -> None:\n        \"\"\"\n        Called in the training loop at the very beginning of the epoch.\n        \"\"\"\n        # do something when the epoch starts\n\n    def on_epoch_end(self) -> None:\n        \"\"\"\n        Called in the training loop at the very end of the epoch.\n        \"\"\"\n        # do something when the epoch ends\n\n    def on_pre_performance_check(self) -> None:\n        \"\"\"\n        Called at the very beginning of the validation loop.\n        \"\"\"\n        # do something before validation starts\n\n    def on_post_performance_check(self) -> None:\n        \"\"\"\n        Called at the very end of the validation loop.\n        \"\"\"\n        # do something before validation end\n\n    def on_before_zero_grad(self, optimizer: Optimizer) -> None:\n        \"\"\"\n        Called after optimizer.step() and before optimizer.zero_grad().\n\n        Called in the training loop after taking an optimizer step and before zeroing grads.\n        Good place to inspect weight information with weights updated.\n\n        This is where it is called::\n\n            for optimizer in optimizers:\n                optimizer.step()\n                model.on_before_zero_grad(optimizer) # < ---- called here\n                optimizer.zero_grad\n\n        Args:\n            optimizer: The optimizer for which grads should be zeroed.\n        \"\"\"\n        # do something with the optimizer or inspect it.\n\n    def on_after_backward(self) -> None:\n        \"\"\"\n        Called in the training loop after loss.backward() and before optimizers do anything.\n        This is the ideal place to inspect or log gradient information.\n\n        Example::\n\n            def on_after_backward(self):\n                # example to inspect gradient information in tensorboard\n                if self.trainer.global_step % 25 == 0:  # don't make the tf file huge\n                    params = self.state_dict()\n                    for k, v in params.items():\n                        grads = v\n                        name = k\n                        self.logger.experiment.add_histogram(tag=name, values=grads,\n                                                             global_step=self.trainer.global_step)\n\n        \"\"\"\n\n    def backward(self, trainer, loss: Tensor, optimizer: Optimizer, optimizer_idx: int) -> None:\n        \"\"\"\n        Override backward with your own implementation if you need to.\n\n        Args:\n            trainer: Pointer to the trainer\n            loss: Loss is already scaled by accumulated grads\n            optimizer: Current optimizer being used\n            optimizer_idx: Index of the current optimizer being used\n\n        Called to perform backward step.\n        Feel free to override as needed.\n\n        The loss passed in has already been scaled for accumulated gradients if requested.\n\n        Example::\n\n            def backward(self, trainer, loss, optimizer, optimizer_idx):\n                loss.backward()\n\n        \"\"\"\n        loss.backward()\n\n    def amp_scale_loss(self, unscaled_loss, optimizer, optimizer_idx):\n        if self.trainer.use_native_amp:\n            scaled_loss = self.trainer.scaler.scale(unscaled_loss)\n\n        else:\n            scaled_loss = amp.scale_loss(unscaled_loss, optimizer)\n\n        return scaled_loss\n\n    def transfer_batch_to_device(self, batch: Any, device: torch.device) -> Any:\n        \"\"\"\n        Override this hook if your :class:`~torch.utils.data.DataLoader` returns tensors\n        wrapped in a custom data structure.\n\n        The data types listed below (and any arbitrary nesting of them) are supported out of the box:\n\n        - :class:`torch.Tensor` or anything that implements `.to(...)`\n        - :class:`list`\n        - :class:`dict`\n        - :class:`tuple`\n        - ``torchtext.data.Batch`` (COMING SOON)\n\n        For anything else, you need to define how the data is moved to the target device (CPU, GPU, TPU, ...).\n\n        Example::\n\n            def transfer_batch_to_device(self, batch, device)\n                if isinstance(batch, CustomBatch):\n                    # move all tensors in your custom data structure to the device\n                    batch.samples = batch.samples.to(device)\n                    batch.targets = batch.targets.to(device)\n                else:\n                    batch = super().transfer_batch_to_device(data, device)\n                return batch\n\n        Args:\n            batch: A batch of data that needs to be transferred to a new device.\n            device: The target device as defined in PyTorch.\n\n        Returns:\n            A reference to the data on the new device.\n\n        Note:\n            This hook should only transfer the data and not modify it, nor should it move the data to\n            any other device than the one passed in as argument (unless you know what you are doing).\n            The :class:`~pytorch_lightning.trainer.trainer.Trainer` already takes care of splitting the\n            batch and determines the target devices.\n\n        See Also:\n            - :func:`~pytorch_lightning.utilities.apply_func.move_data_to_device`\n            - :func:`~pytorch_lightning.utilities.apply_func.apply_to_collection`\n        \"\"\"\n        return move_data_to_device(batch, device)\n", "path": "pytorch_lightning/core/hooks.py"}, {"content": "from abc import ABC\nfrom collections import Mapping, Sequence\nfrom typing import Any, Callable, Union\n\nimport torch\n\n\ndef apply_to_collection(data: Any, dtype: Union[type, tuple], function: Callable, *args, **kwargs) -> Any:\n    \"\"\"\n    Recursively applies a function to all elements of a certain dtype.\n\n    Args:\n        data: the collection to apply the function to\n        dtype: the given function will be applied to all elements of this dtype\n        function: the function to apply\n        *args: positional arguments (will be forwarded to calls of ``function``)\n        **kwargs: keyword arguments (will be forwarded to calls of ``function``)\n\n    Returns:\n        the resulting collection\n\n    \"\"\"\n    elem_type = type(data)\n\n    # Breaking condition\n    if isinstance(data, dtype):\n        return function(data, *args, **kwargs)\n\n    # Recursively apply to collection items\n    elif isinstance(data, Mapping):\n        return elem_type({k: apply_to_collection(v, dtype, function, *args, **kwargs)\n                          for k, v in data.items()})\n    elif isinstance(data, tuple) and hasattr(data, '_fields'):  # named tuple\n        return elem_type(*(apply_to_collection(d, dtype, function, *args, **kwargs) for d in data))\n    elif isinstance(data, Sequence) and not isinstance(data, str):\n        return elem_type([apply_to_collection(d, dtype, function, *args, **kwargs) for d in data])\n\n    # data is neither of dtype, nor a collection\n    return data\n\n\nclass TransferableDataType(ABC):\n    \"\"\"\n    A custom type for data that can be moved to a torch device via `.to(...)`.\n\n    Example:\n\n        >>> isinstance(dict, TransferableDataType)\n        False\n        >>> isinstance(torch.rand(2, 3), TransferableDataType)\n        True\n        >>> class CustomObject:\n        ...     def __init__(self):\n        ...         self.x = torch.rand(2, 2)\n        ...     def to(self, device):\n        ...         self.x = self.x.to(device)\n        ...         return self\n        >>> isinstance(CustomObject(), TransferableDataType)\n        True\n    \"\"\"\n\n    @classmethod\n    def __subclasshook__(cls, subclass):\n        if cls is TransferableDataType:\n            to = getattr(subclass, \"to\", None)\n            return callable(to)\n        return NotImplemented\n\n\ndef move_data_to_device(batch: Any, device: torch.device):\n    \"\"\"\n    Transfers a collection of data to the given device. Any object that defines a method\n    ``to(device)`` will be moved and all other objects in the collection will be left untouched.\n\n    Args:\n        batch: A tensor or collection of tensors or anything that has a method `.to(...)`.\n            See :func:`apply_to_collection` for a list of supported collection types.\n        device: The device to which the data should be moved\n\n    Return:\n        the same collection but with all contained tensors residing on the new device.\n\n    See Also:\n        - :meth:`torch.Tensor.to`\n        - :class:`torch.device`\n    \"\"\"\n    def to(data):\n        return data.to(device, non_blocking=True)\n    return apply_to_collection(batch, dtype=TransferableDataType, function=to)\n", "path": "pytorch_lightning/utilities/apply_func.py"}]}