[DirectML] Fix pow_.

lshqqytiger · Apr 17, 2024 · c113734 · c113734
1 parent 2a61ecc
commit c113734
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/modules/dml/hijack/torch.py b/modules/dml/hijack/torch.py
@@ -45,6 +45,6 @@ def lerp(*args, **kwargs) -> torch.Tensor:
 _pow_ = torch.Tensor.pow_
 def pow_(self: torch.Tensor, *args, **kwargs):
     if self.dtype == torch.float64:
-        return _pow_(self.cpu(), *args, **kwargs)
+        return _pow_(self.cpu(), *args, **kwargs).to(self.device)
     return _pow_(self, *args, **kwargs)
 torch.Tensor.pow_ = _pow_