1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
|
# mypy: allow-untyped-defs
import threading
from typing import Any, Dict
import torch._C._lazy
class DeviceContext:
_CONTEXTS: Dict[str, Any] = {}
_CONTEXTS_LOCK = threading.Lock()
def __init__(self, device):
self.device = device
def get_device_context(device=None):
if device is None:
device = torch._C._lazy._get_default_device_type()
else:
device = str(device)
with DeviceContext._CONTEXTS_LOCK:
devctx = DeviceContext._CONTEXTS.get(device, None)
if devctx is None:
devctx = DeviceContext(device)
DeviceContext._CONTEXTS[device] = devctx
return devctx
|
