fix model on cuda

2019-02-08 11:08:43 +01:00
parent eccb2f0163
commit 5d7e845712
1 changed files with 1 additions and 0 deletions
--- a/examples/run_openai_gpt.py
+++ b/examples/run_openai_gpt.py
@@ -135,6 +135,7 @@ def main():
    tokenizer = OpenAIGPTTokenizer.from_pretrained(args.model_name, special_tokens=special_tokens)
    special_tokens_ids = list(tokenizer.convert_tokens_to_ids(token) for token in special_tokens)
    model = OpenAIGPTDoubleHeadsModel.from_pretrained(args.model_name, num_special_tokens=len(special_tokens))
+    model.to(device)

    # Load and encode the datasets
    def tokenize_and_encode(obj):