Spaces:
Runtime error
Runtime error
| from typing import Optional | |
| import torch | |
| from .utils import rDevice, get_device | |
| class Device: | |
| def __enter__(self, device: Optional[rDevice]=None): | |
| torch.dml.context_device = get_device(device) | |
| def __init__(self, device: Optional[rDevice]=None) -> torch.device: # pylint: disable=return-in-init | |
| return get_device(device) | |
| def __exit__(self, t, v, tb): | |
| torch.dml.context_device = None | |