From d020059f2f71fe7c25765dde9d535195c09ece01 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Sun, 3 Sep 2023 01:14:16 +0200 Subject: Update imports --- text_recognizer/network/vit.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) (limited to 'text_recognizer/network/vit.py') diff --git a/text_recognizer/network/vit.py b/text_recognizer/network/vit.py index b6203d7..1fbf3fc 100644 --- a/text_recognizer/network/vit.py +++ b/text_recognizer/network/vit.py @@ -4,10 +4,10 @@ from typing import Type from einops.layers.torch import Rearrange from torch import Tensor, nn -from text_recognizer.network.transformer.embedding.token import TokenEmbedding -from text_recognizer.network.transformer.embedding.sincos import sincos_2d -from text_recognizer.network.transformer.decoder import Decoder -from text_recognizer.network.transformer.encoder import Encoder +from .transformer.embedding.token import TokenEmbedding +from .transformer.embedding.sincos import sincos_2d +from .transformer.decoder import Decoder +from .transformer.encoder import Encoder class VisionTransformer(nn.Module): @@ -59,11 +59,10 @@ class VisionTransformer(nn.Module): def decode(self, text: Tensor, img_features: Tensor) -> Tensor: text = text.long() - # TODO: add mask to decoder mask = text != self.pad_index tokens = self.token_embedding(text) tokens = tokens + self.pos_embedding(tokens) - output = self.decoder(tokens, context=img_features) + output = self.decoder(tokens, context=img_features, mask=mask) return self.to_logits(output) def forward( -- cgit v1.2.3-70-g09d2