Update rec_train.py

修正from torch import optim
keainye · Sep 10, 2020 · 649339c · 649339c
1 parent f93e47f
commit 649339c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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