update directml

This commit is contained in:
Vladimir Mandic
2023-05-13 11:21:11 -04:00
parent a2485cf7ef
commit d96ab6a1ae
11 changed files with 127 additions and 72 deletions
+4 -4
View File
@@ -6,8 +6,8 @@ import modules.dml.hijack
from .optimizer.unknown import UnknownOptimizer
class DirectML():
def get_optimizer(device: torch.device):
assert(device.type == 'privateuseone')
def get_optimizer(self, device: torch.device):
assert device.type == 'privateuseone'
try:
device_name = torch_directml.device_name(device.index)
if 'NVIDIA' in device_name or 'GeForce' in device_name:
@@ -22,8 +22,8 @@ class DirectML():
except:
return UnknownOptimizer
def memory_stats(device: torch.device):
optimizer = DirectML.get_optimizer(device)
def memory_stats(self, device: torch.device):
optimizer = DirectML.get_optimizer(self, device)
return optimizer.memory_stats(device.index)
# Alternative of torch.cuda for DirectML.