vh/heretic
1
0
Fork 0
mirror of https://github.com/p-e-w/heretic.git synced 2026-09-12 15:16:25 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat(ara): fix issues on some multi-GPU setups

Browse Source 
Co-authored-by: kabachuha <artemkhrapov2001@yandex.ru>
This commit is contained in:
Philipp Emanuel Weidmann
2026-03-07 18:24:31 +05:30
parent 0bb9521fbe
commit d79a443e6f
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/heretic/model.py
+5
View File
@@ -570,6 +570,11 @@ class Model:
module_index module_index
] ]
good_input = good_input.to(matrix.device)
good_output = good_output.to(matrix.device)
bad_input = bad_input.to(matrix.device)
bad_output = bad_output.to(matrix.device)
def objective(matrix: Tensor) -> Tensor: def objective(matrix: Tensor) -> Tensor:
new_good_output = good_input @ matrix.T new_good_output = good_input @ matrix.T
new_bad_output = bad_input @ matrix.T new_bad_output = bad_input @ matrix.T