From 33d814e3e7c91be11517160f423a3a397a72dca8 Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Tue, 25 Jul 2023 00:03:07 +0900 Subject: [PATCH] Fix DirectML tensor behavior. --- modules/dml/hijack/torch.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/modules/dml/hijack/torch.py b/modules/dml/hijack/torch.py index 2400e6f6a..e3d8350df 100644 --- a/modules/dml/hijack/torch.py +++ b/modules/dml/hijack/torch.py @@ -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")