improve device usage

This commit is contained in:
Rémi Louf
2019-12-06 15:45:09 +01:00
committed by Julien Chaumond
parent c0707a85d2
commit 2a64107e44
4 changed files with 21 additions and 19 deletions

View File

@@ -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