Device name should be a string, not bytes (#40322)

malfet · facebook-github-bot · commit 5766da503b7f · 2020-06-22T19:27:25.000-07:00
Summary: I.e. do not accept `bytes` as possible type of `device` argument in `torch.cuda._get_device_index` Pull Request resolved: #40322 Differential Revision: D22176885 Pulled By: malfet fbshipit-source-id: 2f3a46174161f1cdcf6a6ad94a31e54b18ad6186
diff --git a/torch/cuda/_utils.py b/torch/cuda/_utils.py
@@ -1,5 +1,4 @@
 import torch
-import torch._six
 from typing import Optional, Union
 from torch.types import Device
 
@@ -18,7 +17,7 @@ def _get_device_index(device: Union[Device, int], optional=False) -> int:
     If :attr:`device` is ``None``, this will return the current default CUDA
     device if :attr:`optional` is ``True``.
     """
-    if isinstance(device, torch._six.string_classes):
+    if isinstance(device, str):
         device = torch.device(device)
     device_idx: Optional[int]
     if isinstance(device, torch.device):
