3 files changed, 6 insertions, 6 deletions
diff --git a/text_recognizer/models/base.py b/text_recognizer/models/base.py
index 77c5509..886394d 100644
--- a/text_recognizer/models/base.py
+++ b/text_recognizer/models/base.py
@@ -10,7 +10,7 @@ from torch import nn
 from torch import Tensor
 from torchmetrics import Accuracy
 
-from text_recognizer.data.mappings.base import AbstractMapping
+from text_recognizer.data.mappings.base import EmnistMapping
 
 
 class LitBase(LightningModule):
@@ -22,7 +22,7 @@ class LitBase(LightningModule):
         loss_fn: Type[nn.Module],
         optimizer_configs: DictConfig,
         lr_scheduler_configs: Optional[DictConfig],
-        mapping: Type[AbstractMapping],
+        mapping: Type[EmnistMapping],
     ) -> None:
         super().__init__()
 
diff --git a/text_recognizer/models/conformer.py b/text_recognizer/models/conformer.py
index 487eabe..655ebf6 100644
--- a/text_recognizer/models/conformer.py
+++ b/text_recognizer/models/conformer.py
@@ -6,7 +6,7 @@ from omegaconf import DictConfig
 import torch
 from torch import nn, Tensor
 
-from text_recognizer.data.mappings import AbstractMapping
+from text_recognizer.data.mappings import EmnistMapping
 from text_recognizer.models.base import LitBase
 from text_recognizer.models.metrics import CharacterErrorRate
 from text_recognizer.models.util import first_element
@@ -21,7 +21,7 @@ class LitConformer(LitBase):
         loss_fn: Type[nn.Module],
         optimizer_configs: DictConfig,
         lr_scheduler_configs: Optional[DictConfig],
-        mapping: Type[AbstractMapping],
+        mapping: Type[EmnistMapping],
         max_output_len: int = 451,
         start_token: str = "<s>",
         end_token: str = "<e>",
diff --git a/text_recognizer/models/transformer.py b/text_recognizer/models/transformer.py
index 4bbc671..1ffff60 100644
--- a/text_recognizer/models/transformer.py
+++ b/text_recognizer/models/transformer.py
@@ -5,7 +5,7 @@ from omegaconf import DictConfig
 import torch
 from torch import nn, Tensor
 
-from text_recognizer.data.mappings import AbstractMapping
+from text_recognizer.data.mappings import EmnistMapping
 from text_recognizer.models.base import LitBase
 from text_recognizer.models.metrics import CharacterErrorRate
 
@@ -19,7 +19,7 @@ class LitTransformer(LitBase):
         loss_fn: Type[nn.Module],
         optimizer_configs: DictConfig,
         lr_scheduler_configs: Optional[DictConfig],
-        mapping: Type[AbstractMapping],
+        mapping: Type[EmnistMapping],
         max_output_len: int = 451,
         start_token: str = "<s>",
         end_token: str = "<e>",