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

hijack torch.linalg.cholesky()

Browse Source
This commit is contained in:
Seunghoon Lee
2025-09-28 18:43:37 +09:00
parent 9c19e8a9b8
commit 06b14b070b
1 changed files with 10 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/rocm.py
+10 -1
View File
@@ -245,7 +245,7 @@ if sys.platform == "win32":
cholesky_ex_gpu = torch.linalg.cholesky_ex
@wraps(cholesky_ex_gpu)
def cholesky_ex(A: torch.Tensor, upper=False, check_errors=False, out=None) -> torch.Tensor:
def cholesky_ex(A: torch.Tensor, upper=False, check_errors=False, out=None) -> torch.return_types.linalg_cholesky_ex:
if A.device.type != 'cpu':
return cholesky_ex_gpu(A, upper=upper, check_errors=check_errors, out=out)
@@ -269,6 +269,15 @@ if sys.platform == "win32":
return torch.return_types.linalg_cholesky_ex((L, torch.tensor(0, dtype=torch.int32, device='cpu')), {})
torch.linalg.cholesky_ex = cholesky_ex
cholesky_gpu = torch.linalg.cholesky
@wraps(cholesky_gpu)
def cholesky(A: torch.Tensor, upper=False, out=None) -> torch.Tensor:
if A.device.type != 'cpu':
return cholesky_gpu(A, upper=upper, out=out)
L, _ = torch.linalg.cholesky_ex(A, upper=upper, out=out)
return L
torch.linalg.cholesky = cholesky
is_wsl: bool = False
else:
def get_agents() -> List[Agent]: