Restore Python 3.9 compatibility. (DirectML)

This commit is contained in:
Seunghoon Lee
2023-07-29 12:10:41 +09:00
parent 47f2f50574
commit 813eb48bf7
2 changed files with 6 additions and 6 deletions
+4 -4
View File
@@ -1,5 +1,5 @@
# pylint: disable=no-member,no-self-argument,no-method-argument
from typing import Optional
from typing import Optional, Union
import torch
import torch_directml # pylint: disable=import-error
import modules.dml.amp as amp
@@ -13,9 +13,9 @@ class DirectML:
amp = amp
device = device
context_device: torch.device | None = None
context_device: Union[torch.device, None] = None
__gpu_memory_bound: int | None = None
__gpu_memory_bound: Union[int, None] = None
is_autocast_enabled = False
autocast_gpu_dtype = torch.float16
@@ -36,7 +36,7 @@ class DirectML:
except Exception:
return UnknownMemoryControl
def set_gpu_memory_bound(bound: int | None):
def set_gpu_memory_bound(bound: Union[int, None]):
DirectML.__gpu_memory_bound = bound
def is_available() -> bool:
+2 -2
View File
@@ -1,7 +1,7 @@
from typing import Optional
from typing import Optional, Union
import torch
rDevice = torch.device | int
rDevice = Union[torch.device, int]
def get_device(device: Optional[rDevice]=None) -> torch.device:
if device is None:
device = torch.dml.current_device()