mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 17:24:32 +02:00
update directml
This commit is contained in:
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user