Fixes typo in

dgasmith · Sep 6, 2024 · 4fed3e0 · 4fed3e0
1 parent dfb2627
commit 4fed3e0
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/opt_einsum/paths.py b/opt_einsum/paths.py
@@ -500,9 +500,13 @@ def branch(
     size_dict: Dict[str, int],
     memory_limit: Optional[int] = None,
     nbranch: Optional[int] = None,
-    **optimizer_kwargs: Dict[str, Any],
+    cutoff_flops_factor: int = 4,
+    minimize: str = "flops",
+    cost_fn: str = "memory-removed",
 ) -> PathType:
-    optimizer = BranchBound(nbranch=nbranch, **optimizer_kwargs)  # type: ignore
+    optimizer = BranchBound(
+        nbranch=nbranch, cutoff_flops_factor=cutoff_flops_factor, minimize=minimize, cost_fn=cost_fn
+    )
     return optimizer(inputs, output, size_dict, memory_limit)