From 8afa8e1c6e9623b0dea86236da04b2b4173e9443 Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Wed, 7 Apr 2021 22:12:10 +0200
Subject: Fixed typing and typos, train script load config, reformatted

---
 text_recognizer/models/transformer.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

(limited to 'text_recognizer/models/transformer.py')

diff --git a/text_recognizer/models/transformer.py b/text_recognizer/models/transformer.py
index 983e274..b23685b 100644
--- a/text_recognizer/models/transformer.py
+++ b/text_recognizer/models/transformer.py
@@ -1,7 +1,7 @@
 """PyTorch Lightning model for base Transformers."""
 from typing import Dict, List, Optional, Union, Tuple
 
-from omegaconf import OmegaConf
+from omegaconf import DictConfig, OmegaConf
 import pytorch_lightning as pl
 import torch
 from torch import nn
@@ -20,9 +20,9 @@ class LitTransformerModel(LitBaseModel):
     def __init__(
         self,
         network: Type[nn, Module],
-        optimizer: Union[OmegaConf, Dict],
-        lr_scheduler: Union[OmegaConf, Dict],
-        criterion: Union[OmegaConf, Dict],
+        optimizer: Union[DictConfig, Dict],
+        lr_scheduler: Union[DictConfig, Dict],
+        criterion: Union[DictConfig, Dict],
         monitor: str = "val_loss",
         mapping: Optional[List[str]] = None,
     ) -> None:
-- 
cgit v1.2.3-70-g09d2