From 291f4c0dd070f7dcc693cb4233012b6437af1a26 Mon Sep 17 00:00:00 2001 From: wassname Date: Wed, 25 Dec 2019 08:02:01 +0800 Subject: [PATCH] use less GPU ram on startup --- gpt2generator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gpt2generator.py b/gpt2generator.py index 629b841..84b69e6 100644 --- a/gpt2generator.py +++ b/gpt2generator.py @@ -136,7 +136,7 @@ class GPT2Generator: model_class, tokenizer_class = MODEL_CLASSES["gpt2"] self.tokenizer = tokenizer_class.from_pretrained(self.checkpoint_path) self.model = model_class.from_pretrained(self.checkpoint_path) - self.model.to(self.device).to(self.dtype) + self.model.to(self.dtype).to(self.device) self.model.eval() def sample_sequence(self, context_tokens=None, generate_num=None, temperature=None): @@ -193,7 +193,7 @@ class GPT2Generator: truncate_multiple_sequences(context_tokens, self.max_history_tokens) context_tokens = list(itertools.chain(*context_tokens)) - logger.debug("Text passing into model %s", self.tokenizer.decode(o, clean_up_tokenization_spaces=True, skip_special_tokens=True)) + logger.debug("Text passing into model %s", self.tokenizer.decode(context_tokens, clean_up_tokenization_spaces=True, skip_special_tokens=True)) generated = 0 for _ in range(self.samples // self.batch_size):