From d79a443e6fe7fd207dfd2146c9f261e29dfacdd7 Mon Sep 17 00:00:00 2001 From: Philipp Emanuel Weidmann Date: Sat, 7 Mar 2026 18:24:31 +0530 Subject: [PATCH] feat(ara): fix issues on some multi-GPU setups Co-authored-by: kabachuha --- src/heretic/model.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/heretic/model.py b/src/heretic/model.py index 0a4b4fb..898d4df 100644 --- a/src/heretic/model.py +++ b/src/heretic/model.py @@ -570,6 +570,11 @@ class Model: 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: new_good_output = good_input @ matrix.T new_bad_output = bad_input @ matrix.T