Skip to content

Commit

Permalink
[tests] use torch_device instead of auto for model testing (huggi…
Browse files Browse the repository at this point in the history
…ngface#29531)

* use torch_device

* skip for XPU

* Update tests/generation/test_utils.py

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>

---------

Co-authored-by: amyeroberts <22614925+amyeroberts@users.noreply.github.com>
  • Loading branch information
faaany and amyeroberts authored Mar 8, 2024
1 parent 14536c3 commit 1ea3ad1
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions tests/generation/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -1073,6 +1073,9 @@ def test_beam_search_generate_dict_outputs_use_cache(self):
@require_torch_multi_accelerator
def test_model_parallel_beam_search(self):
for model_class in self.all_generative_model_classes:
if "xpu" in torch_device:
return unittest.skip("device_map='auto' does not work with XPU devices")

if model_class._no_split_modules is None:
continue

Expand Down

0 comments on commit 1ea3ad1

Please sign in to comment.