Repository: pytorch/pytorch

Issue:
Improper error msg when calling Variable.cuda() in forked subprocess
In FloatTensorBase: 

RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

In Variable:
RuntimeError: cuda runtime error (3) : initialization error at /py/conda-bld/pytorch_$SOMENUMBER$/work/torch/lib/THC/generic/THCStorage.c:55



Discussion:
Can you please show me the code that gives you that error message?
@apaszke Seems to be more complex than I thought. Need to be .cuda on two different GPUs.

``` python
# from multiprocessing import set_start_method
# try:
#     set_start_method('spawn')
# except RuntimeError:
#     pass

# Above are NOT executed

import torch
from torch.multiprocessing import Process

def boom():
    b = torch.FloatTensor([1,2,3])
    # b = b.cuda(0) # Normal error msg.
    # b = b.cuda(1) # Normal error msg.
    b = torch.autograd.Variable(b)
    # b = b.cuda(1) # Normal error msg.
    b = b.cuda(0) # initialization error at /b/wheel/pytorch-src/torch/lib/THC/generic/THCStorage.c:55

a = torch.FloatTensor([2,2,3])
a = a.cuda(1)
a = torch.autograd.Variable(a)
p = Process(target=boom,args=())
p.daemon = True
p.start()
p.join()

```

Error:
```
THCudaCheck FAIL file=/b/wheel/pytorch-src/torch/lib/THC/generic/THCStorage.c line=55 error=3 : initialization error
Process Process-1:
Traceback (most recent call last):
  File "/usr/lib/python2.7/multiprocessing/process.py", line 258, in _bootstrap
    self.run()
  File "/usr/lib/python2.7/multiprocessing/process.py", line 114, in run
    self._target(*self._args, **self._kwargs)
  File "/------/test_boom.py", line 15, in boom
    b = b.cuda(0) #
  File "/------/torch/autograd/variable.py", line 240, in cuda
    return CudaTransfer(device_id, async)(self)
  File "/------/torch/autograd/_functions/tensor.py", line 160, in forward
    return i.cuda(async=self.async)
  File "/------/torch/_utils.py", line 65, in _cuda
    return new_type(self.size()).copy_(self, async)
RuntimeError: cuda runtime error (3) : initialization error at /b/wheel/pytorch-src/torch/lib/THC/generic/THCStorage.c:55
```
please see http://pytorch.org/docs/notes/multiprocessing.html#sharing-cuda-tensors

This is completely expected, and it is a fundamental CUDA limitation. We published the workaround there.
It's expected that it fails, but it's not expected that it fails with that THC error. I'm reopening this, so we remember to investigate it further

Resolving patch:
diff --git a/torch/cuda/__init__.py b/torch/cuda/__init__.py
--- a/torch/cuda/__init__.py
+++ b/torch/cuda/__init__.py
@@ -96,6 +96,7 @@ def _after_fork(arg):
     if _initialized and _original_pid != os.getpid():
         _initialized = False
         _in_bad_fork = True
+        _CudaBase.__new__ = _lazy_new
 
 
 _register_after_fork(_after_fork, _after_fork)
@@ -260,6 +261,14 @@ def init_err(self):
     torch._C.__dict__['_CudaStreamBase'] = _dummy_type('CudaStreamBase')
 
 
+@staticmethod
+def _lazy_new(cls, *args, **kwargs):
+    _lazy_init()
+    # We need this method only for lazy init, so we can remove it
+    del _CudaBase.__new__
+    return super(_CudaBase, cls).__new__(cls, *args, **kwargs)
+
+
 class _CudaBase(object):
     is_cuda = True
     is_sparse = False
@@ -268,11 +277,7 @@ def type(self, *args, **kwargs):
         with device(self.get_device()):
             return super(_CudaBase, self).type(*args, **kwargs)
 
-    def __new__(cls, *args, **kwargs):
-        _lazy_init()
-        # We need this method only for lazy init, so we can remove it
-        del _CudaBase.__new__
-        return super(_CudaBase, cls).__new__(cls, *args, **kwargs)
+    __new__ = _lazy_new
 
 
 class DoubleStorage(_CudaBase, torch._C.CudaDoubleStorageBase, _StorageBase):


Test patch: