Pytorch_lightning: tensors on wrong device

username: jw3126
cooked: <p>I am trying to use pytorch_lightning with multiple GPU, but get the following error:</p>
<pre><code class="lang-auto">RuntimeError: All input tensors must be on the same device. Received cuda:0 and cuda:3
</code></pre>
<p>How to fix this? Below is a MWE:</p>
<pre><code class="lang-python">import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import pytorch_lightning as pl

class DataModule(pl.LightningDataModule):

    def __init__(self):
        super().__init__()
    
    def setup(self, stage):
        # called on each gpu
        n = 128
        x = torch.randn(n, 1, 64,64)
        data = list(zip(x,x))
        self.test  = DataLoader(data, batch_size=32)
        self.train = DataLoader(data, batch_size=32)
        self.val   = DataLoader(data, batch_size=32)
        
    def train_dataloader(self):
        return self.train

    def val_dataloader(self):
        return self.val

    def test_dataloader(self):
        return self.test
    
class Net(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.net = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=(1,1))
    
    def forward(self, x):
        return self.net(x)
    
    def validation_step(self, batch, batch_idx):
        loss = self.compute_batch_loss(batch, batch_idx)
        self.log('val_loss', loss)
        return loss
    
    def compute_batch_loss(self, batch, batch_idx):
        x, y = batch
        y_hat = self.net(x)
        loss = F.mse_loss(y_hat, y)
        return loss

    def training_step(self, batch, batch_idx):
        loss = self.compute_batch_loss(batch, batch_idx)
        self.log('train_loss', loss)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

dm = DataModule()
model = Net()
trainer = pl.Trainer(gpus=4, 
                     distributed_backend="dp",
                     max_epochs=1,
                    )
trainer.fit(model, dm)
</code></pre>

username: ptrblck
cooked: <p>Your code works fine on my machine and outputs:</p>
<pre><code class="lang-python">Epoch 0: 100%|████████████████████████████████████████████████████| 8/8 [00:00&lt;00:00, 111.96it/s, loss=2.734, v_num=0]
</code></pre>

username: jw3126
cooked: <p>Thanks a lot for digging into this! You also used 4 gpus? Could you rerun with say 10 epochs? This behavior is somewhat random for me and does not always trigger after the first epoch.</p>

username: ptrblck
cooked: <p>I used your provided code on a node with 8 GPUs.<br>
I can try to rerun it later.</p>

username: ptrblck
cooked: <p>Update: I can reproduce the issue using 100 epochs and get:</p>
<pre><code class="lang-python">  File "tmp.py", line 64, in &lt;module&gt;
    trainer.fit(model, dm)
  File "/opt/conda/envs/tmp/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 444, in fit
    results = self.accelerator_backend.train()
  File "/opt/conda/envs/tmp/lib/python3.8/site-packages/pytorch_lightning/accelerators/dp_accelerator.py", line 106, in train
    results = self.train_or_test()
  File "/opt/conda/envs/tmp/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 74, in train_or_test
    results = self.trainer.train()
  File "/opt/conda/envs/tmp/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 493, in train
    self.train_loop.run_training_epoch()
  File "/opt/conda/envs/tmp/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 589, in run_training_epoch
    self.trainer.run_evaluation(test_mode=False)
  File "/opt/conda/envs/tmp/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 609, in run_evaluation
    eval_loop_results = self.evaluation_loop.log_epoch_metrics(deprecated_eval_results, epoch_logs, test_mode)
  File "/opt/conda/envs/tmp/lib/python3.8/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 210, in log_epoch_metrics
    eval_loop_results = self.trainer.logger_connector.on_evaluation_epoch_end(
  File "/opt/conda/envs/tmp/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/logger_connector.py", line 113, in on_evaluation_epoch_end
    self._log_on_evaluation_epoch_end_metrics(epoch_logs)
  File "/opt/conda/envs/tmp/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/logger_connector.py", line 181, in _log_on_evaluation_epoch_end_metrics
    reduced_epoch_metrics = dl_metrics[0].__class__.reduce_on_epoch_end(dl_metrics)
  File "/opt/conda/envs/tmp/lib/python3.8/site-packages/pytorch_lightning/core/step_result.py", line 464, in reduce_on_epoch_end
    recursive_stack(result)
  File "/opt/conda/envs/tmp/lib/python3.8/site-packages/pytorch_lightning/core/step_result.py", line 603, in recursive_stack
    result[k] = collate_tensors(v)
  File "/opt/conda/envs/tmp/lib/python3.8/site-packages/pytorch_lightning/core/step_result.py", line 625, in collate_tensors
    return torch.stack(items)
RuntimeError: All input tensors must be on the same device. Received cuda:1 and cuda:3
</code></pre>
<p>I’m not familiar enough with PyTorch Lightning and would suggest to create an issue with this code snippet in their repository.</p>
<p>CC <a class="mention" href="/u/williamfalcon">@williamFalcon</a> for visibility.</p>

username: jw3126
cooked: <p>Thanks a lot for reproducing!</p>

username: jw3126
cooked: <aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/PyTorchLightning/pytorch-lightning/issues/4423" target="_blank" rel="noopener nofollow ugc">github.com/PyTorchLightning/pytorch-lightning</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/PyTorchLightning/pytorch-lightning/issues/4423" target="_blank" rel="noopener nofollow ugc">Weird behaviour multi-GPU (dp, gpus &gt; 1) tensors not converted to cuda</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-10-29" data-time="09:26:32" data-timezone="UTC">09:26AM - 29 Oct 20 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/ClaartjeBarkhof" target="_blank" rel="noopener nofollow ugc">
          <img alt="ClaartjeBarkhof" src="https://avatars3.githubusercontent.com/u/25668035?v=4" class="onebox-avatar-inline" width="20" height="20">
          ClaartjeBarkhof
        </a>
      </div>
    </div>
  </div>
</div>

<div class="github-row">
  <p class="github-content">🐛 Bug
I encounter unexpected behaviour for different versions of Pytorch Lightning:

When using PL version 1.0.4 my train step is ignored all...</p>
</div>

<div class="labels">
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">DP</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">bug / fix</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">help wanted</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">waiting on author</span>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>