Match the returned # of arguments/gradients between CPU hungarian mat…

…ching custom VJP primal function and bwd rule. PiperOrigin-RevId: 719440273
google-research · Jan 24, 2025 · 8820fc9 · 8820fc9
1 parent 794bb40
commit 8820fc9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/scenic/model_lib/matchers/common.py b/scenic/model_lib/matchers/common.py
@@ -110,7 +110,7 @@ def matching_fn_hcb_vjp_fwd(cost, n_cols):
     return matching_fn_hcb(cost, n_cols), None
 
   def matching_fn_hcb_vjp_bwd(*_):
-    return (None,)  # Return no gradient.
+    return (None, None)  # Return no gradient.
 
   matching_fn_hcb.defvjp(matching_fn_hcb_vjp_fwd, matching_fn_hcb_vjp_bwd)