Repository: pytorch
Issue: Simple C++ custom autograd function code throws error "CUDA error: driver shutting down"
Labels: module: cpp,module: autograd,triaged

## 🐛 Bug

Running the following code in cuda-enabled libtorch throws error "CUDA error: driver shutting down", even though the code doesn't use CUDA. Running the same code in cpu-only libtorch doesn't throw any error.

```cpp
#include <iostream>
#include <torch/torch.h>

using namespace torch::autograd;

class MulConstant : public Function<MulConstant> {
 public:
  static Variable forward(AutogradContext *ctx, Variable variable, double constant) {
    ctx->saved_data["constant"] = constant;
    return variable * constant;
  }

  static variable_list backward(AutogradContext *ctx, variable_list grad_outputs) {
    return {grad_outputs[0] * ctx->saved_data["constant"].toDouble(), Variable()};
  }
};

int main(int argc, char* argv[])
{
  auto x = torch::randn({2}).requires_grad_();
  auto y = MulConstant::apply(x, 5.5);
  y.sum().backward();
  std::cout << x.grad() << std::endl;
}
```
Error:
```
terminate called after throwing an instance of 'c10::Error'
terminate called recursively
  what():  CUDA error: driver shutting down (setDevice at /pytorch/c10/cuda/impl/CUDAGuardImpl.h:42)
frame #0: c10::Error::Error(c10::SourceLocation, std::string const&) + 0x46 (0x7fedc6342656 in /data/libtorch/libtorch_nightly_cu92/libtorch/lib/libc10.so)
frame #1: <unknown function> + 0xc6c2 (0x7fed7bad26c2 in /data/libtorch/libtorch_nightly_cu92/libtorch/lib/libc10_cuda.so)
frame #2: torch::autograd::Engine::set_device(int) + 0x159 (0x7fedb9c36b39 in /data/libtorch/libtorch_nightly_cu92/libtorch/lib/libtorch_cpu.so)
frame #3: torch::autograd::Engine::thread_init(int, std::shared_ptr<torch::autograd::ReadyQueue> const&) + 0x34 (0x7fedb9c39064 in /data/libtorch/libtorch_nightly_cu92/libtorch/lib/libtorch_cpu.so)
frame #4: <unknown function> + 0xc70f (0x7fedc657b70f in /data/libtorch/libtorch_nightly_cu92/libtorch/lib/libtorch.so)
frame #5: <unknown function> + 0x76ba (0x7fed7c3756ba in /lib/x86_64-linux-gnu/libpthread.so.0)
frame #6: clone + 0x6d (0x7fed7c8bc41d in /lib/x86_64-linux-gnu/libc.so.6)
Aborted (core dumped)
```
Better backtrace:
```
Thread 4 "example-app" hit Catchpoint 1 (exception thrown), 0x00007fffccab38bd in __cxa_throw () from /usr/lib/x86_64-linux-gnu/libstdc++.so.6
(gdb) bt
#0  0x00007fffccab38bd in __cxa_throw () from /usr/lib/x86_64-linux-gnu/libstdc++.so.6
#1  0x00007fffc4d14ab9 in c10::cuda::impl::CUDAGuardImpl::getDevice (this=0x997920) at ../c10/cuda/impl/CUDAGuardImpl.h:37
#2  0x00007fffc4d14ed6 in c10::cuda::impl::CUDAGuardImpl::setDevice (this=0x997920, d=...) at ../c10/cuda/impl/CUDAGuardImpl.h:51
#3  0x00007ffff0f101db in torch::autograd::Engine::set_device (this=0x7ffff7b16bc0 <torch::autograd::Engine::get_base_engine()::engine>, device=1) at ../torch/csrc/autograd/engine.cpp:264
#4  0x00007ffff0f1034d in torch::autograd::Engine::thread_init (this=0x7ffff7b16bc0 <torch::autograd::Engine::get_base_engine()::engine>, device=1, ready_queue=std::shared_ptr (count 2, weak 0) 0x1b33aa0)
    at ../torch/csrc/autograd/engine.cpp:293
#5  0x00007ffff0f3613e in std::_Mem_fn_base<void (torch::autograd::Engine::*)(int, std::shared_ptr<torch::autograd::ReadyQueue> const&), true>::operator()<int, std::shared_ptr<torch::autograd::ReadyQueue>, void>(torch::autograd::Engine*, int&&, std::shared_ptr<torch::autograd::ReadyQueue>&&) const (this=0x1b340d8, __object=0x7ffff7b16bc0 <torch::autograd::Engine::get_base_engine()::engine>) at /usr/include/c++/5/functional:600
#6  0x00007ffff0f360a1 in std::_Bind_simple<std::_Mem_fn<void (torch::autograd::Engine::*)(int, std::shared_ptr<torch::autograd::ReadyQueue> const&)> (torch::autograd::Engine*, int, std::shared_ptr<torch::autograd::ReadyQueue>)>::_M_invoke<0ul, 1ul, 2ul>(std::_Index_tuple<0ul, 1ul, 2ul>) (this=0x1b340b8) at /usr/include/c++/5/functional:1531
#7  0x00007ffff0f35cb8 in std::_Bind_simple<std::_Mem_fn<void (torch::autograd::Engine::*)(int, std::shared_ptr<torch::autograd::ReadyQueue> const&)> (torch::autograd::Engine*, int, std::shared_ptr<torch::autograd::ReadyQueue>)>::operator()() (this=0x1b340b8) at /usr/include/c++/5/functional:1520
#8  0x00007ffff0f35ac8 in std::thread::_Impl<std::_Bind_simple<std::_Mem_fn<void (torch::autograd::Engine::*)(int, std::shared_ptr<torch::autograd::ReadyQueue> const&)> (torch::autograd::Engine*, int, std::shared_ptr<torch::autograd::ReadyQueue>)> >::_M_run() (this=0x1b340a0) at /usr/include/c++/5/thread:115
#9  0x00007fffccadec80 in ?? () from /usr/lib/x86_64-linux-gnu/libstdc++.so.6
#10 0x00007fffcbafa6ba in start_thread (arg=0x7fff9cd8d700) at pthread_create.c:333
#11 0x00007fffcc24441d in clone () at ../sysdeps/unix/sysv/linux/x86_64/clone.S:109
```
Update: I noticed that if we initialize a cuda tensor (e.g. `auto cuda_tensor = torch::randn({3, 4}, torch::kCUDA); std::cout << cuda_tensor << std::endl;`) before running the C++ custom autograd function, the whole thing would pass and there is no error.

## Expected behavior

It should just work without throwing any error.

## Environment

Latest libtorch nightly

cc @ezyang @SsnL @albanD @zou3519 @gqchen @yf225