From af9b31766a66242c8846c365a22851c2559d67ef Mon Sep 17 00:00:00 2001
From: DanEnergetics <d.mann95.dm@gmail.com>
Date: Mon, 26 Aug 2024 15:25:09 +0200
Subject: [PATCH] update typing

Co-authored-by: michelwi <michelwi@users.noreply.github.com>
---
 i6_models/parts/fsa.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/i6_models/parts/fsa.py b/i6_models/parts/fsa.py
index 18944235..aef723a8 100644
--- a/i6_models/parts/fsa.py
+++ b/i6_models/parts/fsa.py
@@ -83,7 +83,7 @@ def build_single(self, seq_tag: str) -> Tuple[int, int, np.ndarray, np.ndarray]:
         raw_fsa = self.builder.build_by_segment_name(seq_tag)
         return raw_fsa
 
-    def build_batch(self, seq_tags: Iterable[str]) -> TWeightedFsa:
+    def build_batch(self, seq_tags: Iterable[str]) -> WeightedFsa:
         """
         Build and concatenate the FSAs for a batch of sequence tags
         and reformat as an input to `i6_native_ops.fbw.fbw_loss`.