From aa5738c18e2e43cef6d66497a19841068fb49833 Mon Sep 17 00:00:00 2001
From: Franz Srambical <79149449+emergenz@users.noreply.github.com>
Date: Thu, 15 May 2025 10:28:01 +0200
Subject: [PATCH] fix: wrong type signature

---
 utils/nn.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/utils/nn.py b/utils/nn.py
index 2c6db160..021e2dc7 100644
--- a/utils/nn.py
+++ b/utils/nn.py
@@ -1,5 +1,5 @@
 import math
-from typing import Dict
+from typing import Dict, Tuple
 
 from flax import linen as nn
 import jax
@@ -112,7 +112,7 @@ def setup(self):
         )
         self.drop = nn.Dropout(self.dropout, deterministic=False)
 
-    def __call__(self, x: jax.Array, training: bool) -> Dict[str, jax.Array]:
+    def __call__(self, x: jax.Array, training: bool) -> Tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
         # --- Compute distances ---
         x = normalize(x)
         codebook = normalize(self.codebook)