Update delta-iris/src/tokenizer.py

delta-iris/src/tokenizer.py

@@ -8,7 +8,7 @@ import torch.nn as nn
 
 from .models.convnet import FrameEncoder, FrameDecoder
 from .data import Batch
-from .models.tokenizer.quantizer import Quantizer, QuantizerOutput
+from .models.tokenizer.quantizer import Quantizer
 from .models.utils import init_weights, LossWithIntermediateLosses
 
 class Tokenizer(nn.Module):
@@ -38,7 +38,7 @@ class Tokenizer(nn.Module):
     def __repr__(self) -> str:
         return "tokenizer"
 
-    def forward(self, x1: torch.FloatTensor, a: torch.LongTensor, x2: torch.FloatTensor) -> QuantizerOutput:
+    def forward(self, x1: torch.FloatTensor, a: torch.LongTensor, x2: torch.FloatTensor) -> dict:
         z = self.encode(x1, a, x2)
         z = rearrange(z, 'b t c (h k) (w l) -> b t (h w) (k l c)', h=self.tokens_grid_res, w=self.tokens_grid_res)