Update delta-iris/src/models/tokenizer/quantizer.py

delta-iris/src/models/tokenizer/quantizer.py

@@ -7,15 +7,6 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-
-@dataclass
-class QuantizerOutput:
-    q: torch.FloatTensor
-    tokens: torch.LongTensor
-    loss: Dict[str, torch.FloatTensor]
-    metrics: Dict[str, float]
-
-
 class Quantizer(nn.Module):
     def __init__(self, codebook_size: int, codebook_dim: int, input_dim: int, max_codebook_updates_with_revival: Optional[int] = None) -> None:
         super().__init__()
@@ -51,8 +42,12 @@ class Quantizer(nn.Module):
 
         q = rearrange(q, '(b t k) e -> b t k e', b=b, k=k)
         tokens = rearrange(tokens, '(b t k) -> b t k', b=b, k=k)
-
-        return QuantizerOutput(q, tokens, losses, metrics)
+        return {
+            "q": q, 
+            "tokens": tokens,
+            "losses": losses, 
+            "metrics": metrics
+        }
 
     @torch.no_grad()
     def update_codebook(self, z: torch.Tensor, tokens: torch.LongTensor) -> None: