improve device usage

This commit is contained in:
Rémi Louf authored and Julien Chaumond committed 2019-12-09 20:37:55 -05:00
1 parent c0707a85d2
commit 2a64107e44
4 files changed
+21 -19

No files matched your search

@@ -847,14 +847,12 @@ class Translator(object):
global_scores (:obj:`GlobalScorer`):
object to rescore final translations
copy_attn (bool): use copy attention during translation
cuda (bool): use cuda
beam_trace (bool): trace beam search for debugging
logger(logging.Logger): logger.
"""
def __init__(self, args, model, vocab, symbols, global_scorer=None, logger=None):
self.logger = logger
self.cuda = args.visible_gpus != "-1"
self.args = args
self.model = model