Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-20 01:31:13 +02:00
Code Issues Packages Projects Releases Wiki Activity

Fix DirectML tensor behavior.

Browse Source
This commit is contained in:
Seunghoon Lee
2023-07-25 00:03:07 +09:00
parent 14eec801ec
commit 33d814e3e7
1 changed files with 14 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/dml/hijack/torch.py
+14 -7
View File
@@ -2,12 +2,19 @@ import torch
from modules.sd_hijack_utils import CondFunc
def to_sub(orig, self: torch.Tensor, *args, **kwargs):
def validate(device: torch.device | str):
if torch.dml.is_directml_device(torch.device(device)):
raise NotImplementedError("Cannot copy out of meta tensor; no data!")
for arg in args:
validate(arg)
if "device" in kwargs:
validate(kwargs["device"])
return orig(self, *args, **kwargs)
CondFunc('torchsde._brownian.brownian_interval._randn', lambda _, size, dtype, device, seed: torch.randn(size, dtype=dtype, device=torch.device("cpu"), generator=torch.Generator(torch.device("cpu")).manual_seed(int(seed))).to(device), lambda _, size, dtype, device, seed: device.type == 'privateuseone')
_new = torch.Tensor.new
def new(self: torch.Tensor, *args, **kwargs):
if torch.dml.is_directml_device(self.device):
return _new(self.cpu(), *args, **kwargs).to(self.device)
return _new(self, *args, **kwargs)
torch.Tensor.new = new
# https://github.com/microsoft/DirectML/issues/400
CondFunc('torch.Tensor.new', lambda orig, self, *args, **kwargs: orig(self.cpu(), *args, **kwargs), lambda orig, self, *args, **kwargs: torch.dml.is_directml_device(self.device))
# https://github.com/microsoft/DirectML/issues/477
CondFunc('torch.Tensor.to', to_sub, lambda orig, self, *args, **kwargs: self.device.type == "meta")