fix generic error w genot

theislab · Nov 25, 2024 · 0fcd620 · 0fcd620
1 parent baa0362
commit 0fcd620
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/moscot/problems/generic/_generic.py b/src/moscot/problems/generic/_generic.py
@@ -790,7 +790,7 @@ def prepare(
         cost: OttCostFn_t = "sq_euclidean",
         cost_kwargs: CostKwargs_t = types.MappingProxyType({}),
         **kwargs: Any,
-    ) -> "GENOTLinProblem[K, B]":
+    ) -> "GENOTLinProblem":
         """Prepare the :class:`moscot.problems.generic.GENOTLinProblem`."""
         self.batch_key = key
         xy, kwargs = handle_joint_attr_tmp(joint_attr, kwargs)
@@ -816,7 +816,7 @@ def solve(
         valid_sinkhorn_kwargs: Dict[str, Any] = MappingProxyType({}),
         train_size: float = 1.0,
         **kwargs: Any,
-    ) -> "GENOTLinProblem[K, B]":
+    ) -> "GENOTLinProblem":
         """Solve."""
         return super().solve(
             batch_size=batch_size,