Lazy init in set device, also should not be called in getDevCount

csarofeen · csarofeen · commit 41e480ac0f4d · 2018-01-29T14:44:40.000-08:00
diff --git a/torch/cuda/__init__.py b/torch/cuda/__init__.py
@@ -202,10 +202,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:
@@ -291,7 +291,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
