From 649339c3690c0a805defc8282bb96e0b3ca834a1 Mon Sep 17 00:00:00 2001 From: light1003 <zhanlong.hao@gmail.com> Date: Thu, 10 Sep 2020 13:59:23 +0800 Subject: [PATCH] Update rec_train.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 修正from torch import optim --- tools/rec_train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/rec_train.py b/tools/rec_train.py index c02b7e4..0427d50 100644 --- a/tools/rec_train.py +++ b/tools/rec_train.py @@ -20,7 +20,7 @@ import torch from tqdm import tqdm from torch import nn - +from torch import optim from torchocr.networks import build_model, build_loss from torchocr.datasets import build_dataloader from torchocr.utils import get_logger, weight_init, load_checkpoint, save_checkpoint @@ -77,7 +77,7 @@ def build_optimizer(params, config): 优化器 Returns: """ - from torch import optim + opt_type = config.pop('type') opt = getattr(optim, opt_type)(params, **config) return opt