fix encoding wrapper for batchencoding

webis-de · Dec 10, 2024 · 8cdd245 · 8cdd245
1 parent 56f2f08
commit 8cdd245
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/lightning_ir/base/model.py b/lightning_ir/base/model.py
@@ -230,7 +230,10 @@ def _cat_outputs(
         for key, value in output.items():
             agg[key].append(value)
             types[key] = type(value)
-    return OutputClass(**{key: _cat_outputs(value, types[key]) for key, value in agg.items()})
+    kwargs = {key: _cat_outputs(value, types[key]) for key, value in agg.items()}
+    if OutputClass is BatchEncoding:
+        return OutputClass(kwargs)
+    return OutputClass(**kwargs)
 
 
 class BatchEncodingWrapper(Protocol):