Tracing callback hotfix (#1234)

* fixed bug in TracingCallback

* removed unnecessary imports

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)