Repository: catalyst-team/catalyst

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



Discussion:
🤔  @elephantmipt why do we have hardcoded `cpu` version?
@y-ksenia have you tried the same callback without this line?
Just checked - works fine without cpu transfering.

Resolving patch:
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 @@ def on_stage_end(self, runner: "IRunner") -> None:
         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)
 


Test patch: