Expected all tensors to be on the same device

Nobody has claimed this yet.

### Bug description

I try to implement neural machine translition model using PyTorch Lightning (also known as pl) for training process. While we use pl, we can't use .to('cuda') or .cuda() functions for tensors. When I train my model without pl, the whole process went fine, but after using pl I have this error (use 2 gpu in kaggle):

```
Cell In[70], line 21, in EncoderTransformerLayer.forward(self, value, key, query, mask)
18 def forward(self, value, key, query, mask):
19 # attn_output = self.dropout(self.norm(self.attention(value, key, query, mask)))
20 # mlp_output = self.dropout(self.norm(self.mlp(attn_output))) # сюда pre ln
---> 21 value = self.norm_for_v(value)
22 key = self.norm_for_k(key)
23 query = self.norm_for_q(query)

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument weight in method wrapper_CUDA__native_layer_norm)
```

### What version are you seeing the problem on?

v2.0

### How to reproduce the bug

```python
The output shows that the error is in the following code:

class EncoderTransformerLayer(pl.LightningModule):
def __init__(self, hidden_dim: int, num_heads: int, dropout: float = 0.1):
super().__init__()

self.attention = AttentionModule(hidden_dim, num_heads)
self.mlp = MLP(hidden_dim)

# self.norm = torch.nn.LayerNorm(hidden_dim)
# self.dropout = torch.nn.Dropout(dropout)
self.norm_for_v = torch.nn.LayerNorm(hidden_dim)
self.norm_for_k = torch.nn.LayerNorm(hidden_dim)
self.norm_for_q = torch.nn.LayerNorm(hidden_dim)

self.norm_for_attention = torch.nn.LayerNorm(hidden_dim)
self.norm_for_mlp = torch.nn.LayerNorm(hidden_dim)

def forward(self, value, key, query, mask):
# attn_output = self.dropout(self.norm(self.attention(value, key, query, mask)))
# mlp_output = self.dropout(self.norm(self.mlp(attn_output))) # сюда pre ln
value = self.norm_for_v(value)
key = self.norm_for_k(key)
query = self.norm_for_q(query)

attn_output = self.attention(value, key, query, mask)
attn_output = self.norm_for_attention(attn_output)

mlp_output = self.mlp(attn_output)
mlp_output = self.norm_for_mlp(mlp_output)

return mlp_output
```

Full bug report:
```
RuntimeError Traceback (most recent call last)
Cell In[78], line 6
1 trainer = pl.Trainer(accelerator="gpu",
2 devices=1,
3 max_epochs=10,
4 precision='16-mixed')
----> 6 trainer.fit(model=model,
7 train_dataloaders=train_dataloader,
8 val_dataloaders=valid_dataloader,
9 )

File /opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:520, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
518 model = _maybe_unwrap_optimized(model)
519 self.strategy._lightning_module = model
--> 520 call._call_and_handle_interrupt(
521 self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
522 )

File /opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py:44, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)
42 return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
43 else:
---> 44 return trainer_fn(*args, **kwargs)
46 except _TunerExitException:
47 _call_teardown_hook(trainer)

File /opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:559, in Trainer._fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
549 self._data_connector.attach_data(
550 model, train_dataloaders=train_dataloaders, val_dataloaders=val_dataloaders, datamodule=datamodule
551 )
553 ckpt_path = self._checkpoint_connector._select_ckpt_path(
554 self.state.fn,
555 ckpt_path,
556 model_provided=True,
557 model_connected=self.lightning_module is not None,
558 )
--> 559 self._run(model, ckpt_path=ckpt_path)
561 assert self.state.stopped
562 self.training = False

File /opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:935, in Trainer._run(self, model, ckpt_path)
930 self._signal_connector.register_signal_handlers()
932 # ----------------------------
933 # RUN THE TRAINER
934 # ----------------------------
--> 935 results = self._run_stage()
937 # ----------------------------
938 # POST-Training CLEAN UP
939 # ----------------------------
940 log.debug(f"{self.__class__.__name__}: trainer tearing down")

File /opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:976, in Trainer._run_stage(self)
974 if self.training:
975 with isolate_rng():
--> 976 self._run_sanity_check()
977 with torch.autograd.set_detect_anomaly(self._detect_anomaly):
978 self.fit_loop.run()

File /opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1005, in Trainer._run_sanity_check(self)
1002 call._call_callback_hooks(self, "on_sanity_check_start")
1004 # run eval step
-> 1005 val_loop.run()
1007 call._call_callback_hooks(self, "on_sanity_check_end")
1009 # reset logger connector

File /opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/utilities.py:177, in _no_grad_context.._decorator(self, *args, **kwargs)
175 context_manager = torch.no_grad
176 with context_manager():
--> 177 return loop_run(self, *args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py:115, in _EvaluationLoop.run(self)
113 previous_dataloader_idx = dataloader_idx
114 # run step hooks
--> 115 self._evaluation_step(batch, batch_idx, dataloader_idx)
116 except StopIteration:
117 # this needs to wrap the `*_step` call too (not just `next`) for `dataloader_iter` support
118 break

File /opt/conda/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py:375, in _EvaluationLoop._evaluation_step(self, batch, batch_idx, dataloader_idx)
372 self.batch_progress.increment_started()
374 hook_name = "test_step" if trainer.testing else "validation_step"
--> 375 output = call._call_strategy_hook(trainer, hook_name, *step_kwargs.values())
377 self.batch_progress.increment_processed()
379 hook_name = "on_test_batch_end" if trainer.testing else "on_validation_batch_end"

File /opt/conda/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py:288, in _call_strategy_hook(trainer, hook_name, *args, **kwargs)
285 return
287 with trainer.profiler.profile(f"[Strategy]{trainer.strategy.__class__.__name__}.{hook_name}"):
--> 288 output = fn(*args, **kwargs)
290 # restore current_fx when nested context
291 pl_module._current_fx_name = prev_fx_name

File /opt/conda/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py:378, in Strategy.validation_step(self, *args, **kwargs)
376 with self.precision_plugin.val_step_context():
377 assert isinstance(self.model, ValidationStep)
--> 378 return self.model.validation_step(*args, **kwargs)

Cell In[75], line 75, in TranslationModel.validation_step(self, batch, batch_idx)
73 src_mask = self.make_src_mask(src_ids)
74 trg_mask = self.make_trg_mask(trg_ids)
---> 75 encoder_output = self.encoder(src_ids, src_mask)
76 decoder_output = self.decoder(trg_ids, encoder_output, src_mask, trg_mask)
77 trg = trg[:, 1:].contiguous().view(-1)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
1496 # If we don't have any hooks, we want to skip the rest of the logic in
1497 # this function, and just call forward.
1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1499 or _global_backward_pre_hooks or _global_backward_hooks
1500 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501 return forward_call(*args, **kwargs)
1502 # Do not call functions when jit is used
1503 full_backward_hooks, non_full_backward_hooks = [], []

Cell In[71], line 34, in Encoder.forward(self, inputs, mask)
32 #тут был to.deivce
33 for layer in self.layers:
---> 34 hidden_dim = layer(hidden_dim, hidden_dim, hidden_dim, mask)
36 return hidden_dim

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
1496 # If we don't have any hooks, we want to skip the rest of the logic in
1497 # this function, and just call forward.
1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1499 or _global_backward_pre_hooks or _global_backward_hooks
1500 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501 return forward_call(*args, **kwargs)
1502 # Do not call functions when jit is used
1503 full_backward_hooks, non_full_backward_hooks = [], []

Cell In[70], line 21, in EncoderTransformerLayer.forward(self, value, key, query, mask)
18 def forward(self, value, key, query, mask):
19 # attn_output = self.dropout(self.norm(self.attention(value, key, query, mask))) # сюда pre ln
20 # mlp_output = self.dropout(self.norm(self.mlp(attn_output))) # сюда pre ln
---> 21 value = self.norm_for_v(value)
22 key = self.norm_for_k(key)
23 query = self.norm_for_q(query)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
1496 # If we don't have any hooks, we want to skip the rest of the logic in
1497 # this function, and just call forward.
1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1499 or _global_backward_pre_hooks or _global_backward_hooks
1500 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501 return forward_call(*args, **kwargs)
1502 # Do not call functions when jit is used
1503 full_backward_hooks, non_full_backward_hooks = [], []

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/normalization.py:190, in LayerNorm.forward(self, input)
189 def forward(self, input: Tensor) -> Tensor:
--> 190 return F.layer_norm(
191 input, self.normalized_shape, self.weight, self.bias, self.eps)

File /opt/conda/lib/python3.10/site-packages/torch/nn/functional.py:2515, in layer_norm(input, normalized_shape, weight, bias, eps)
2511 if has_torch_function_variadic(input, weight, bias):
2512 return handle_torch_function(
2513 layer_norm, (input, weight, bias), input, normalized_shape, weight=weight, bias=bias, eps=eps
2514 )
-> 2515 return torch.layer_norm(input, normalized_shape, weight, bias, eps, torch.backends.cudnn.enabled)

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument weight in method wrapper_CUDA__native_layer_norm)

```

```

### Environment

Current environment

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):
```

### More info

_No response_

Start with the reported Trainer.fit call and notebook cells 70, 71, 75, and 78, especially EncoderTransformerLayer.forward and the validation path. Reproduce the failure with the supplied model and environment details, then determine whether it is caused by Lightning or the user model. Done requires a confirmed repository-level reproduction and a clearly defined expected behavior.

Written by the indexing model from the issue text.