Lazy init in set device, also should not be called in getDevCount (#4918)

csarofeen · soumith · commit ef4cf860accd · 2018-01-30T16:24:31.000+01:00
diff --git a/torch/cuda/__init__.py b/torch/cuda/__init__.py
@@ -215,10 +215,10 @@ def __init__(self, idx):
     def __enter__(self):
         if self.idx is -1:
             return
-        _lazy_init()
         self.prev_idx = torch._C._cuda_getDevice()
         if self.prev_idx != self.idx:
             torch._C._cuda_setDevice(self.idx)
+        _lazy_init()
 
     def __exit__(self, *args):
         if self.prev_idx != self.idx:
@@ -308,7 +308,6 @@ def stream(stream):
 def device_count():
     """Returns the number of GPUs available."""
     if is_available():
-        _lazy_init()
         return torch._C._cuda_getDeviceCount()
     else:
         return 0
