Repository: pytorch
Issue: Backward function causes device error in C++ when changing module's device repeatly.
Labels: module: cpp,module: autograd,triaged

# 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

Steps to reproduce the behavior:

1. Compile & Run code below.

```C++
#include <torch/torch.h>

// Define a new Module.
struct Net : torch::nn::Module {
  Net() {
    // Construct and register two Linear submodules.
    fc1 = register_module("fc1", torch::nn::Linear(784, 64));
    fc2 = register_module("fc2", torch::nn::Linear(64, 32));
    fc3 = register_module("fc3", torch::nn::Linear(32, 10));
  }

  // Implement the Net's algorithm.
  torch::Tensor forward(torch::Tensor x) {
    // Use one of many tensor manipulation functions.
    x = torch::relu(fc1->forward(x.reshape({x.size(0), 784})));
    x = torch::dropout(x, /*p=*/0.5, /*train=*/is_training());
    x = torch::relu(fc2->forward(x));
    x = torch::log_softmax(fc3->forward(x), /*dim=*/1);
    return x;
  }

  // Use one of many "standard library" modules.
  torch::nn::Linear fc1{nullptr}, fc2{nullptr}, fc3{nullptr};
};

int main() {
  // Create a new Net.
  auto net = std::make_shared<Net>();

  // Create a multi-threaded data loader for the MNIST dataset.
  auto data_loader = torch::data::make_data_loader(
      torch::data::datasets::MNIST("./data").map(
          torch::data::transforms::Stack<>()),
      /*batch_size=*/64);

  // Instantiate an SGD optimization algorithm to update our Net's parameters.
  torch::optim::SGD optimizer(net->parameters(), /*lr=*/0.01);

  for (size_t epoch = 1; epoch <= 10; ++epoch) {
    size_t batch_index = 0;
    // Iterate the data loader to yield batches from the dataset.
    torch::Device device = epoch % 2 == 1 ? torch::kCUDA : torch::kCPU;
    net->to(device);
    std::cout << "Device of network is " << device << std::endl;
    for (auto& batch : *data_loader) {
      // Reset gradients.
      optimizer.zero_grad();
      // Execute the model on the input data.
      torch::Tensor prediction = net->forward(batch.data.to(device));
      // Compute a loss value to judge the prediction of our model.
      torch::Tensor loss = torch::nll_loss(prediction, batch.target.to(device));
      // Compute gradients of the loss w.r.t. the parameters of our model.
      if(batch_index == 0) std::cout << "Device of loss is " << loss.device() << "\n";
      loss.backward();
      if(batch_index == 0) std::cout << "loss.backward() success" << "\n";
	 	
      // Update the parameters based on the calculated gradients.
      optimizer.step();
      // Output the loss and checkpoint every 100 batches.
      if (++batch_index % 100 == 0) {
        std::cout << "Epoch: " << epoch << " | Batch: " << batch_index
                  << " | Loss: " << loss.item<float>() << std::endl;
        // Serialize your model periodically as a checkpoint.
        torch::save(net, "net.pt");
      }
    }
  }
}
```

and output is

```
Device of network is cuda
Device of loss is cuda:0
loss.backward() success
Epoch: 1 | Batch: 100 | Loss: 2.2953
Epoch: 1 | Batch: 200 | Loss: 2.252
Epoch: 1 | Batch: 300 | Loss: 2.22181
Epoch: 1 | Batch: 400 | Loss: 2.13327
Epoch: 1 | Batch: 500 | Loss: 2.04376
Epoch: 1 | Batch: 600 | Loss: 1.80017
Epoch: 1 | Batch: 700 | Loss: 1.65806
Epoch: 1 | Batch: 800 | Loss: 1.47708
Epoch: 1 | Batch: 900 | Loss: 1.40919
Device of network is cpu
Device of loss is cpu
terminate called after throwing an instance of 'torch::utils::FutureError'
  what():  expected device cuda:0 but got device cpu (compute_types at /pytorch/aten/src/ATen/native/TensorIterator.cpp:246)
frame #0: c10::Error::Error(c10::SourceLocation, std::string const&) + 0x46 (0x7f0850510576 in /home/sonic/Downloaded_Library/libtorch/lib/libc10.so)
frame #1: at::TensorIterator::compute_types() + 0x17d4 (0x7f08426fe214 in /home/sonic/Downloaded_Library/libtorch/lib/libtorch_cpu.so)
frame #2: at::TensorIterator::build() + 0x44 (0x7f0842700104 in /home/sonic/Downloaded_Library/libtorch/lib/libtorch_cpu.so)
frame #3: at::TensorIterator::binary_op(at::Tensor&, at::Tensor const&, at::Tensor const&, bool) + 0x146 (0x7f08427007b6 in /home/sonic/Downloaded_Library/libtorch/lib/libtorch_cpu.so)
frame #4: at::native::add_out(at::Tensor&, at::Tensor const&, at::Tensor const&, c10::Scalar) + 0x3c (0x7f084242155c in /home/sonic/Downloaded_Library/libtorch/lib/libtorch_cpu.so)
frame #5: <unknown function> + 0x101717d (0x7f0805af717d in /home/sonic/Downloaded_Library/libtorch/lib/libtorch_cuda.so)
frame #6: <unknown function> + 0x2b517ea (0x7f084443e7ea in /home/sonic/Downloaded_Library/libtorch/lib/libtorch_cpu.so)
frame #7: <unknown function> + 0x2cf9c5c (0x7f08445e6c5c in /home/sonic/Downloaded_Library/libtorch/lib/libtorch_cpu.so)
frame #8: torch::autograd::AccumulateGrad::apply(std::vector<at::Tensor, std::allocator<at::Tensor> >&&) + 0xdc (0x7f08445e81bc in /home/sonic/Downloaded_Library/libtorch/lib/libtorch_cpu.so)
frame #9: <unknown function> + 0x2cf5745 (0x7f08445e2745 in /home/sonic/Downloaded_Library/libtorch/lib/libtorch_cpu.so)
frame #10: torch::autograd::Engine::evaluate_function(std::shared_ptr<torch::autograd::GraphTask>&, torch::autograd::Node*, torch::autograd::InputBuffer&) + 0x16f3 (0x7f08445dfa43 in /home/sonic/Downloaded_Library/libtorch/lib/libtorch_cpu.so)
frame #11: torch::autograd::Engine::thread_main(std::shared_ptr<torch::autograd::GraphTask> const&, bool) + 0x3d2 (0x7f08445e0822 in /home/sonic/Downloaded_Library/libtorch/lib/libtorch_cpu.so)
frame #12: torch::autograd::Engine::thread_init(int) + 0x39 (0x7f08445d8e99 in /home/sonic/Downloaded_Library/libtorch/lib/libtorch_cpu.so)
frame #13: <unknown function> + 0xc70f (0x7f085074970f in /home/sonic/Downloaded_Library/libtorch/lib/libtorch.so)
frame #14: <unknown function> + 0x76db (0x7f08416d56db in /lib/x86_64-linux-gnu/libpthread.so.0)
frame #15: clone + 0x3f (0x7f0803dec88f in /lib/x86_64-linux-gnu/libc.so.6)

Aborted (core dumped)
```

## Expected behavior

Run w/o any segmentation faults.

## Environment

PyTorch version: 1.2.0
Is debug build: No
CUDA used to build PyTorch: 10.0.130

OS: Ubuntu 18.04.4 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: version 3.10.2

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.0.130
GPU models and configuration: GPU 0: GeForce RTX 2070 SUPER
Nvidia driver version: 430.50
cuDNN version: /usr/local/cuda-10.0/targets/x86_64-linux/lib/libcudnn.so.7.5.0

Versions of relevant libraries:
[pip3] numpy==1.17.1
[pip3] torch==1.2.0
[pip3] torchvision==0.4.0
[conda] Could not collect

cc @ezyang @SsnL @albanD @zou3519 @gqchen @yf225