From 0fcd620dba413fd28f0a1d1f046c720f13045174 Mon Sep 17 00:00:00 2001
From: selmanozleyen <syozleyen@gmail.com>
Date: Mon, 25 Nov 2024 16:20:14 +0100
Subject: [PATCH] fix generic error w genot

---
 src/moscot/problems/generic/_generic.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/moscot/problems/generic/_generic.py b/src/moscot/problems/generic/_generic.py
index d18ab794..fcd3d8b2 100644
--- 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,