From 7e0a0a39a54fd7d1a69b9f12bbd98a2b16285c9c Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Mon, 11 Oct 2021 22:09:51 +0200 Subject: Add mapping to base lit model --- text_recognizer/models/base.py | 4 ++++ 1 file changed, 4 insertions(+) (limited to 'text_recognizer') diff --git a/text_recognizer/models/base.py b/text_recognizer/models/base.py index 34f40a2..8aadc39 100644 --- a/text_recognizer/models/base.py +++ b/text_recognizer/models/base.py @@ -11,6 +11,7 @@ from torch import nn from torch import Tensor import torchmetrics +from text_recognizer.data.mappings.base_mapping import AbstractMapping @attr.s(eq=False) class BaseLitModel(LightningModule): @@ -24,6 +25,9 @@ class BaseLitModel(LightningModule): loss_fn: Type[nn.Module] = attr.ib() optimizer_configs: DictConfig = attr.ib() lr_scheduler_configs: Optional[DictConfig] = attr.ib() + mapping: Type[AbstractMapping] = attr.ib() + + # Placeholders train_acc: torchmetrics.Accuracy = attr.ib( init=False, default=torchmetrics.Accuracy() ) -- cgit v1.2.3-70-g09d2