Fixed device bug that occurred when merging best info in TRUNCATED ba…

…ckward modes. (#181)
facebookresearch · Apr 26, 2022 · 733deff · 733deff
1 parent 4ae3d03
commit 733deff
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/theseus/optimizer/nonlinear/nonlinear_optimizer.py b/theseus/optimizer/nonlinear/nonlinear_optimizer.py
@@ -182,7 +182,7 @@ def _merge_infos(
             best_solution = {}
             best_err_no_grad = info.best_err
             best_err_grad = grad_loop_info.best_err
-            idx_no_grad = best_err_no_grad < best_err_grad
+            idx_no_grad = (best_err_no_grad < best_err_grad).cpu().view(-1, 1)
             best_err = torch.minimum(best_err_no_grad, best_err_grad)
             for var_name in info.best_solution:
                 sol_no_grad = info.best_solution[var_name]