Initialize cuda before setting cuda tensor types as default

zou3519 · apaszke · commit 8c69eacde671 · 2018-01-23T11:06:22.000+01:00
diff --git a/torch/__init__.py b/torch/__init__.py
@@ -129,6 +129,11 @@ def set_default_tensor_type(t):
     global Storage
     Tensor = _import_dotted_name(t)
     Storage = _import_dotted_name(t.replace('Tensor', 'Storage'))
+
+    if 'cuda' in t:
+        import torch.cuda
+        torch.cuda.init()
+
     _C._set_default_tensor_type(Tensor)
 
 
