Commit ff22b3ac authored by Stanislas Polu's avatar Stanislas Polu
Browse files

Few small nits in GPT-2's code examples

parent cbb7fad3
...@@ -427,7 +427,7 @@ model.to('cuda') ...@@ -427,7 +427,7 @@ model.to('cuda')
with torch.no_grad(): with torch.no_grad():
hidden_states_1, past = model(tokens_tensor_1) hidden_states_1, past = model(tokens_tensor_1)
# past can be used to reuse precomputed hidden state in a subsequent predictions # past can be used to reuse precomputed hidden state in a subsequent predictions
# (see beam-search examples in the run_gpt2.py example # (see beam-search examples in the run_gpt2.py example).
hidden_states_2, past = model(tokens_tensor_2, past=past) hidden_states_2, past = model(tokens_tensor_2, past=past)
``` ```
...@@ -439,14 +439,15 @@ model = GPT2LMHeadModel.from_pretrained('gpt2') ...@@ -439,14 +439,15 @@ model = GPT2LMHeadModel.from_pretrained('gpt2')
model.eval() model.eval()
# If you have a GPU, put everything on cuda # If you have a GPU, put everything on cuda
tokens_tensor_1 = tokens_tensor.to('cuda') tokens_tensor_1 = tokens_tensor_1.to('cuda')
tokens_tensor_2 = tokens_tensor_2.to('cuda')
model.to('cuda') model.to('cuda')
# Predict all tokens # Predict all tokens
with torch.no_grad(): with torch.no_grad():
predictions_1, past = model(tokens_tensor_1) predictions_1, past = model(tokens_tensor_1)
# past can be used to reuse precomputed hidden state in a subsequent predictions # past can be used to reuse precomputed hidden state in a subsequent predictions
# (see beam-search examples in the run_gpt2.py example # (see beam-search examples in the run_gpt2.py example).
predictions_2, past = model(tokens_tensor_2, past=past) predictions_2, past = model(tokens_tensor_2, past=past)
# get the predicted last token # get the predicted last token
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment