Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions tests/generation/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3037,7 +3037,7 @@ def test_eos_token_id_int_and_list_greedy_search(self):
text = """Hello, my dog is cute and"""
tokens = tokenizer(text, return_tensors="pt")

model = GPT2LMHeadModel.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel it's better to keep it is, but change (line above)

tokens = tokenizer(text, return_tensors="pt")

to

tokens = tokenizer(text, return_tensors="pt").to(torch_device)

WDYT?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TIL that can be done :D

Changed!

model = GPT2LMHeadModel.from_pretrained("hf-internal-testing/tiny-random-gpt2")

torch.manual_seed(0)
eos_token_id = 873
Expand All @@ -3062,7 +3062,7 @@ def test_eos_token_id_int_and_list_contrastive_search(self):
text = """Hello, my dog is cute and"""
tokens = tokenizer(text, return_tensors="pt")

model = GPT2LMHeadModel.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
model = GPT2LMHeadModel.from_pretrained("hf-internal-testing/tiny-random-gpt2")

torch.manual_seed(0)
eos_token_id = 225
Expand All @@ -3088,7 +3088,7 @@ def test_eos_token_id_int_and_list_top_k_top_sampling(self):
text = """Hello, my dog is cute and"""
tokens = tokenizer(text, return_tensors="pt")

model = GPT2LMHeadModel.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
model = GPT2LMHeadModel.from_pretrained("hf-internal-testing/tiny-random-gpt2")

torch.manual_seed(0)
eos_token_id = 846
Expand All @@ -3111,7 +3111,7 @@ def test_eos_token_id_int_and_list_beam_search(self):
text = """Hello, my dog is cute and"""
tokens = tokenizer(text, return_tensors="pt")

model = GPT2LMHeadModel.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(torch_device)
model = GPT2LMHeadModel.from_pretrained("hf-internal-testing/tiny-random-gpt2")

torch.manual_seed(0)
eos_token_id = 873
Expand Down