We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
TracingCallback fails in ConfigAPI
## 🐛 Bug Report
<!-- A clear and concise description of what the bug is. -->
When running TracingCallback the run fails with an RuntimeError.

Apparently, inside the callback, there is a [transfer from any device to `cpu`](https://github.com/catalyst-team/catalyst/blob/a927bed68c1abcd3d28811d1d3d9d7b95e0861a1/catalyst/callbacks/tracing.py#L150) of batch which is used for tracing. And when the `runner.model` is on `cuda` it fails.

### How To Reproduce
Steps to reproduce the behavior:
1. Add 
```
      trace:
        _target_: TracingCallback
        input_key: *model_input
        logdir: *logdir
```
to [mnist_stages config](https://github.com/catalyst-team/catalyst/blob/master/examples/mnist_stages/config.yml).
2. `sh mnist_stages/run_config.sh`

<!-- If you have a code sample, error messages, stack traces, please provide it here as well --> 

### Expected behavior
<!-- A clear and concise description of what you expected to happen. -->
`RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same`


### Environment
```bash
# example checklist, fill with your info
Catalyst version: 21.05
PyTorch version: 1.7.1
Is debug build: No
CUDA used to build PyTorch: 10.2
TensorFlow version: N/A
TensorBoard version: N/A

OS: Linux Mint 19.3 Tricia
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: version 3.10.2

Python version: 3.8
Is CUDA available: Yes
CUDA runtime version: 9.1.85
GPU models and configuration:
GPU 0: GeForce RTX 2070 SUPER
GPU 1: Quadro P4000

Nvidia driver version: 460.32.03
cuDNN version: Could not collect
```


### Additional context
<!-- Add any other context about the problem here. -->


### Checklist
- [x] bug description
- [x] steps to reproduce
- [x] expected behavior
- [x] environment
- [ ] code sample / screenshots
  
### FAQ
Please review the FAQ before submitting an issue:
- [ ] I have read the [documentation and FAQ](https://catalyst-team.github.io/catalyst/)
- [ ] I have reviewed the [minimal examples section](https://github.com/catalyst-team/catalyst#minimal-examples)
- [ ] I have checked the [changelog](https://github.com/catalyst-team/catalyst/blob/master/CHANGELOG.md) for main framework updates
- [ ] I have read the [contribution guide](https://github.com/catalyst-team/catalyst/blob/master/CONTRIBUTING.md)
- [ ] I have joined [Catalyst slack (#__questions channel)](https://join.slack.com/t/catalyst-team-core/shared_invite/zt-d9miirnn-z86oKDzFMKlMG4fgFdZafw) for issue discussion


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `catalyst/callbacks/tracing.py`
Content:
```
1 from typing import List, TYPE_CHECKING, Union
2 from pathlib import Path
3 
4 import torch
5 
6 from catalyst.core import Callback, CallbackNode, CallbackOrder
7 from catalyst.utils.torch import any2device
8 from catalyst.utils.tracing import trace_model
9 
10 if TYPE_CHECKING:
11     from catalyst.core import IRunner
12 
13 
14 class TracingCallback(Callback):
15     """
16     Callback for model tracing.
17 
18     Args:
19         input_key: input key from ``runner.batch`` to use for model tracing
20         logdir: path to folder for saving
21         filename: filename
22         method_name: Model's method name that will be used as entrypoint during tracing
23 
24     Example:
25 
26         .. code-block:: python
27 
28             import os
29 
30             import torch
31             from torch import nn
32             from torch.utils.data import DataLoader
33 
34             from catalyst import dl
35             from catalyst.data import ToTensor
36             from catalyst.contrib.datasets import MNIST
37             from catalyst.contrib.nn.modules import Flatten
38 
39             loaders = {
40                 "train": DataLoader(
41                     MNIST(
42                         os.getcwd(), train=False, download=True, transform=ToTensor()
43                     ),
44                     batch_size=32,
45                 ),
46                 "valid": DataLoader(
47                     MNIST(
48                         os.getcwd(), train=False, download=True, transform=ToTensor()
49                     ),
50                     batch_size=32,
51                 ),
52             }
53 
54             model = nn.Sequential(
55                 Flatten(), nn.Linear(784, 512), nn.ReLU(), nn.Linear(512, 10)
56             )
57             criterion = nn.CrossEntropyLoss()
58             optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
59             runner = dl.SupervisedRunner()
60             runner.train(
61                 model=model,
62                 callbacks=[dl.TracingCallback(input_key="features", logdir="./logs")],
63                 loaders=loaders,
64                 criterion=criterion,
65                 optimizer=optimizer,
66                 num_epochs=1,
67                 logdir="./logs",
68             )
69     """
70 
71     def __init__(
72         self,
73         input_key: Union[str, List[str]],
74         logdir: Union[str, Path] = None,
75         filename: str = "traced_model.pth",
76         method_name: str = "forward",
77     ):
78         """
79         Callback for model tracing.
80 
81         Args:
82             input_key: input key from ``runner.batch`` to use for model tracing
83             logdir: path to folder for saving
84             filename: filename
85             method_name: Model's method name that will be used as entrypoint during tracing
86 
87         Example:
88             .. code-block:: python
89 
90                 import os
91 
92                 import torch
93                 from torch import nn
94                 from torch.utils.data import DataLoader
95 
96                 from catalyst import dl
97                 from catalyst.data import ToTensor
98                 from catalyst.contrib.datasets import MNIST
99                 from catalyst.contrib.nn.modules import Flatten
100 
101                 loaders = {
102                     "train": DataLoader(
103                         MNIST(
104                             os.getcwd(), train=False, download=True, transform=ToTensor()
105                         ),
106                         batch_size=32,
107                     ),
108                     "valid": DataLoader(
109                         MNIST(
110                             os.getcwd(), train=False, download=True, transform=ToTensor()
111                         ),
112                         batch_size=32,
113                     ),
114                 }
115 
116                 model = nn.Sequential(
117                     Flatten(), nn.Linear(784, 512), nn.ReLU(), nn.Linear(512, 10)
118                 )
119                 criterion = nn.CrossEntropyLoss()
120                 optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
121                 runner = dl.SupervisedRunner()
122                 runner.train(
123                     model=model,
124                     callbacks=[dl.TracingCallback(input_key="features", logdir="./logs")],
125                     loaders=loaders,
126                     criterion=criterion,
127                     optimizer=optimizer,
128                     num_epochs=1,
129                     logdir="./logs",
130                 )
131         """
132         super().__init__(order=CallbackOrder.ExternalExtra, node=CallbackNode.Master)
133         if logdir is not None:
134             self.filename = str(Path(logdir) / filename)
135         else:
136             self.filename = filename
137         self.method_name = method_name
138 
139         self.input_key = [input_key] if isinstance(input_key, str) else input_key
140 
141     def on_stage_end(self, runner: "IRunner") -> None:
142         """
143         On stage end action.
144 
145         Args:
146             runner: runner for experiment
147         """
148         model = runner.model
149         batch = tuple(runner.batch[key] for key in self.input_key)
150         batch = any2device(batch, "cpu")
151         traced_model = trace_model(model=model, batch=batch, method_name=self.method_name)
152         torch.jit.save(traced_model, self.filename)
153 
154 
155 __all__ = ["TracingCallback"]
156
```
--- 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/catalyst/callbacks/tracing.py b/catalyst/callbacks/tracing.py
--- a/catalyst/callbacks/tracing.py
+++ b/catalyst/callbacks/tracing.py
@@ -4,7 +4,6 @@
 import torch
 
 from catalyst.core import Callback, CallbackNode, CallbackOrder
-from catalyst.utils.torch import any2device
 from catalyst.utils.tracing import trace_model
 
 if TYPE_CHECKING:
@@ -145,9 +144,9 @@
         Args:
             runner: runner for experiment
         """
-        model = runner.model
+        model = runner.engine.sync_device(runner.model)
         batch = tuple(runner.batch[key] for key in self.input_key)
-        batch = any2device(batch, "cpu")
+        batch = runner.engine.sync_device(batch)
         traced_model = trace_model(model=model, batch=batch, method_name=self.method_name)
         torch.jit.save(traced_model, self.filename)

{"golden_diff": "diff --git a/catalyst/callbacks/tracing.py b/catalyst/callbacks/tracing.py\n--- a/catalyst/callbacks/tracing.py\n+++ b/catalyst/callbacks/tracing.py\n@@ -4,7 +4,6 @@\n import torch\n \n from catalyst.core import Callback, CallbackNode, CallbackOrder\n-from catalyst.utils.torch import any2device\n from catalyst.utils.tracing import trace_model\n \n if TYPE_CHECKING:\n@@ -145,9 +144,9 @@\n         Args:\n             runner: runner for experiment\n         \"\"\"\n-        model = runner.model\n+        model = runner.engine.sync_device(runner.model)\n         batch = tuple(runner.batch[key] for key in self.input_key)\n-        batch = any2device(batch, \"cpu\")\n+        batch = runner.engine.sync_device(batch)\n         traced_model = trace_model(model=model, batch=batch, method_name=self.method_name)\n         torch.jit.save(traced_model, self.filename)\n", "issue": "TracingCallback fails in ConfigAPI\n## \ud83d\udc1b Bug Report\r\n<!-- A clear and concise description of what the bug is. -->\r\nWhen running TracingCallback the run fails with an RuntimeError.\r\n\r\nApparently, inside the callback, there is a [transfer from any device to `cpu`](https://github.com/catalyst-team/catalyst/blob/a927bed68c1abcd3d28811d1d3d9d7b95e0861a1/catalyst/callbacks/tracing.py#L150) of batch which is used for tracing. And when the `runner.model` is on `cuda` it fails.\r\n\r\n### How To Reproduce\r\nSteps to reproduce the behavior:\r\n1. Add \r\n```\r\n      trace:\r\n        _target_: TracingCallback\r\n        input_key: *model_input\r\n        logdir: *logdir\r\n```\r\nto [mnist_stages config](https://github.com/catalyst-team/catalyst/blob/master/examples/mnist_stages/config.yml).\r\n2. `sh mnist_stages/run_config.sh`\r\n\r\n<!-- If you have a code sample, error messages, stack traces, please provide it here as well --> \r\n\r\n### Expected behavior\r\n<!-- A clear and concise description of what you expected to happen. -->\r\n`RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same`\r\n\r\n\r\n### Environment\r\n```bash\r\n# example checklist, fill with your info\r\nCatalyst version: 21.05\r\nPyTorch version: 1.7.1\r\nIs debug build: No\r\nCUDA used to build PyTorch: 10.2\r\nTensorFlow version: N/A\r\nTensorBoard version: N/A\r\n\r\nOS: Linux Mint 19.3 Tricia\r\nGCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0\r\nCMake version: version 3.10.2\r\n\r\nPython version: 3.8\r\nIs CUDA available: Yes\r\nCUDA runtime version: 9.1.85\r\nGPU models and configuration:\r\nGPU 0: GeForce RTX 2070 SUPER\r\nGPU 1: Quadro P4000\r\n\r\nNvidia driver version: 460.32.03\r\ncuDNN version: Could not collect\r\n```\r\n\r\n\r\n### Additional context\r\n<!-- Add any other context about the problem here. -->\r\n\r\n\r\n### Checklist\r\n- [x] bug description\r\n- [x] steps to reproduce\r\n- [x] expected behavior\r\n- [x] environment\r\n- [ ] code sample / screenshots\r\n  \r\n### FAQ\r\nPlease review the FAQ before submitting an issue:\r\n- [ ] I have read the [documentation and FAQ](https://catalyst-team.github.io/catalyst/)\r\n- [ ] I have reviewed the [minimal examples section](https://github.com/catalyst-team/catalyst#minimal-examples)\r\n- [ ] I have checked the [changelog](https://github.com/catalyst-team/catalyst/blob/master/CHANGELOG.md) for main framework updates\r\n- [ ] I have read the [contribution guide](https://github.com/catalyst-team/catalyst/blob/master/CONTRIBUTING.md)\r\n- [ ] I have joined [Catalyst slack (#__questions channel)](https://join.slack.com/t/catalyst-team-core/shared_invite/zt-d9miirnn-z86oKDzFMKlMG4fgFdZafw) for issue discussion\r\n\n", "before_files": [{"content": "from typing import List, TYPE_CHECKING, Union\nfrom pathlib import Path\n\nimport torch\n\nfrom catalyst.core import Callback, CallbackNode, CallbackOrder\nfrom catalyst.utils.torch import any2device\nfrom catalyst.utils.tracing import trace_model\n\nif TYPE_CHECKING:\n    from catalyst.core import IRunner\n\n\nclass TracingCallback(Callback):\n    \"\"\"\n    Callback for model tracing.\n\n    Args:\n        input_key: input key from ``runner.batch`` to use for model tracing\n        logdir: path to folder for saving\n        filename: filename\n        method_name: Model's method name that will be used as entrypoint during tracing\n\n    Example:\n\n        .. code-block:: python\n\n            import os\n\n            import torch\n            from torch import nn\n            from torch.utils.data import DataLoader\n\n            from catalyst import dl\n            from catalyst.data import ToTensor\n            from catalyst.contrib.datasets import MNIST\n            from catalyst.contrib.nn.modules import Flatten\n\n            loaders = {\n                \"train\": DataLoader(\n                    MNIST(\n                        os.getcwd(), train=False, download=True, transform=ToTensor()\n                    ),\n                    batch_size=32,\n                ),\n                \"valid\": DataLoader(\n                    MNIST(\n                        os.getcwd(), train=False, download=True, transform=ToTensor()\n                    ),\n                    batch_size=32,\n                ),\n            }\n\n            model = nn.Sequential(\n                Flatten(), nn.Linear(784, 512), nn.ReLU(), nn.Linear(512, 10)\n            )\n            criterion = nn.CrossEntropyLoss()\n            optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)\n            runner = dl.SupervisedRunner()\n            runner.train(\n                model=model,\n                callbacks=[dl.TracingCallback(input_key=\"features\", logdir=\"./logs\")],\n                loaders=loaders,\n                criterion=criterion,\n                optimizer=optimizer,\n                num_epochs=1,\n                logdir=\"./logs\",\n            )\n    \"\"\"\n\n    def __init__(\n        self,\n        input_key: Union[str, List[str]],\n        logdir: Union[str, Path] = None,\n        filename: str = \"traced_model.pth\",\n        method_name: str = \"forward\",\n    ):\n        \"\"\"\n        Callback for model tracing.\n\n        Args:\n            input_key: input key from ``runner.batch`` to use for model tracing\n            logdir: path to folder for saving\n            filename: filename\n            method_name: Model's method name that will be used as entrypoint during tracing\n\n        Example:\n            .. code-block:: python\n\n                import os\n\n                import torch\n                from torch import nn\n                from torch.utils.data import DataLoader\n\n                from catalyst import dl\n                from catalyst.data import ToTensor\n                from catalyst.contrib.datasets import MNIST\n                from catalyst.contrib.nn.modules import Flatten\n\n                loaders = {\n                    \"train\": DataLoader(\n                        MNIST(\n                            os.getcwd(), train=False, download=True, transform=ToTensor()\n                        ),\n                        batch_size=32,\n                    ),\n                    \"valid\": DataLoader(\n                        MNIST(\n                            os.getcwd(), train=False, download=True, transform=ToTensor()\n                        ),\n                        batch_size=32,\n                    ),\n                }\n\n                model = nn.Sequential(\n                    Flatten(), nn.Linear(784, 512), nn.ReLU(), nn.Linear(512, 10)\n                )\n                criterion = nn.CrossEntropyLoss()\n                optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)\n                runner = dl.SupervisedRunner()\n                runner.train(\n                    model=model,\n                    callbacks=[dl.TracingCallback(input_key=\"features\", logdir=\"./logs\")],\n                    loaders=loaders,\n                    criterion=criterion,\n                    optimizer=optimizer,\n                    num_epochs=1,\n                    logdir=\"./logs\",\n                )\n        \"\"\"\n        super().__init__(order=CallbackOrder.ExternalExtra, node=CallbackNode.Master)\n        if logdir is not None:\n            self.filename = str(Path(logdir) / filename)\n        else:\n            self.filename = filename\n        self.method_name = method_name\n\n        self.input_key = [input_key] if isinstance(input_key, str) else input_key\n\n    def on_stage_end(self, runner: \"IRunner\") -> None:\n        \"\"\"\n        On stage end action.\n\n        Args:\n            runner: runner for experiment\n        \"\"\"\n        model = runner.model\n        batch = tuple(runner.batch[key] for key in self.input_key)\n        batch = any2device(batch, \"cpu\")\n        traced_model = trace_model(model=model, batch=batch, method_name=self.method_name)\n        torch.jit.save(traced_model, self.filename)\n\n\n__all__ = [\"TracingCallback\"]\n", "path": "catalyst/callbacks/tracing.py"}], "after_files": [{"content": "from typing import List, TYPE_CHECKING, Union\nfrom pathlib import Path\n\nimport torch\n\nfrom catalyst.core import Callback, CallbackNode, CallbackOrder\nfrom catalyst.utils.tracing import trace_model\n\nif TYPE_CHECKING:\n    from catalyst.core import IRunner\n\n\nclass TracingCallback(Callback):\n    \"\"\"\n    Callback for model tracing.\n\n    Args:\n        input_key: input key from ``runner.batch`` to use for model tracing\n        logdir: path to folder for saving\n        filename: filename\n        method_name: Model's method name that will be used as entrypoint during tracing\n\n    Example:\n\n        .. code-block:: python\n\n            import os\n\n            import torch\n            from torch import nn\n            from torch.utils.data import DataLoader\n\n            from catalyst import dl\n            from catalyst.data import ToTensor\n            from catalyst.contrib.datasets import MNIST\n            from catalyst.contrib.nn.modules import Flatten\n\n            loaders = {\n                \"train\": DataLoader(\n                    MNIST(\n                        os.getcwd(), train=False, download=True, transform=ToTensor()\n                    ),\n                    batch_size=32,\n                ),\n                \"valid\": DataLoader(\n                    MNIST(\n                        os.getcwd(), train=False, download=True, transform=ToTensor()\n                    ),\n                    batch_size=32,\n                ),\n            }\n\n            model = nn.Sequential(\n                Flatten(), nn.Linear(784, 512), nn.ReLU(), nn.Linear(512, 10)\n            )\n            criterion = nn.CrossEntropyLoss()\n            optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)\n            runner = dl.SupervisedRunner()\n            runner.train(\n                model=model,\n                callbacks=[dl.TracingCallback(input_key=\"features\", logdir=\"./logs\")],\n                loaders=loaders,\n                criterion=criterion,\n                optimizer=optimizer,\n                num_epochs=1,\n                logdir=\"./logs\",\n            )\n    \"\"\"\n\n    def __init__(\n        self,\n        input_key: Union[str, List[str]],\n        logdir: Union[str, Path] = None,\n        filename: str = \"traced_model.pth\",\n        method_name: str = \"forward\",\n    ):\n        \"\"\"\n        Callback for model tracing.\n\n        Args:\n            input_key: input key from ``runner.batch`` to use for model tracing\n            logdir: path to folder for saving\n            filename: filename\n            method_name: Model's method name that will be used as entrypoint during tracing\n\n        Example:\n            .. code-block:: python\n\n                import os\n\n                import torch\n                from torch import nn\n                from torch.utils.data import DataLoader\n\n                from catalyst import dl\n                from catalyst.data import ToTensor\n                from catalyst.contrib.datasets import MNIST\n                from catalyst.contrib.nn.modules import Flatten\n\n                loaders = {\n                    \"train\": DataLoader(\n                        MNIST(\n                            os.getcwd(), train=False, download=True, transform=ToTensor()\n                        ),\n                        batch_size=32,\n                    ),\n                    \"valid\": DataLoader(\n                        MNIST(\n                            os.getcwd(), train=False, download=True, transform=ToTensor()\n                        ),\n                        batch_size=32,\n                    ),\n                }\n\n                model = nn.Sequential(\n                    Flatten(), nn.Linear(784, 512), nn.ReLU(), nn.Linear(512, 10)\n                )\n                criterion = nn.CrossEntropyLoss()\n                optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)\n                runner = dl.SupervisedRunner()\n                runner.train(\n                    model=model,\n                    callbacks=[dl.TracingCallback(input_key=\"features\", logdir=\"./logs\")],\n                    loaders=loaders,\n                    criterion=criterion,\n                    optimizer=optimizer,\n                    num_epochs=1,\n                    logdir=\"./logs\",\n                )\n        \"\"\"\n        super().__init__(order=CallbackOrder.ExternalExtra, node=CallbackNode.Master)\n        if logdir is not None:\n            self.filename = str(Path(logdir) / filename)\n        else:\n            self.filename = filename\n        self.method_name = method_name\n\n        self.input_key = [input_key] if isinstance(input_key, str) else input_key\n\n    def on_stage_end(self, runner: \"IRunner\") -> None:\n        \"\"\"\n        On stage end action.\n\n        Args:\n            runner: runner for experiment\n        \"\"\"\n        model = runner.engine.sync_device(runner.model)\n        batch = tuple(runner.batch[key] for key in self.input_key)\n        batch = runner.engine.sync_device(batch)\n        traced_model = trace_model(model=model, batch=batch, method_name=self.method_name)\n        torch.jit.save(traced_model, self.filename)\n\n\n__all__ = [\"TracingCallback\"]\n", "path": "catalyst/callbacks/tracing.py"}]}