From bd4bd443f339e95007bfdabf3e060db720f4d4b9 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Tue, 3 Aug 2021 18:18:48 +0200 Subject: Training working, multiple bug fixes --- text_recognizer/data/base_data_module.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) (limited to 'text_recognizer/data/base_data_module.py') diff --git a/text_recognizer/data/base_data_module.py b/text_recognizer/data/base_data_module.py index fd914b6..16a06d9 100644 --- a/text_recognizer/data/base_data_module.py +++ b/text_recognizer/data/base_data_module.py @@ -1,12 +1,12 @@ """Base lightning DataModule class.""" from pathlib import Path -from typing import Dict, Tuple +from typing import Dict, Tuple, Type import attr from pytorch_lightning import LightningDataModule from torch.utils.data import DataLoader -from text_recognizer.data.mappings import AbstractMapping +from text_recognizer.data.base_mapping import AbstractMapping from text_recognizer.data.base_dataset import BaseDataset @@ -25,7 +25,7 @@ class BaseDataModule(LightningDataModule): def __attrs_pre_init__(self) -> None: super().__init__() - mapping: AbstractMapping = attr.ib() + mapping: Type[AbstractMapping] = attr.ib() batch_size: int = attr.ib(default=16) num_workers: int = attr.ib(default=0) pin_memory: bool = attr.ib(default=True) -- cgit v1.2.3-70-g09d2