[tests] use torch_device instead of auto for model testing

[faaany]

What does this PR do?

When running test cases under the models folder on XPU, I found that many model tests fail at the same test test_model_parallel_beam_search, e.g.

FAILED tests/models/bigbird_pegasus/test_modeling_bigbird_pegasus.py::BigBirdPegasusModelTest::test_model_parallel_beam_search - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and xpu:0! (when che

This is because device_map="auto" is used. As elaborated in this PR, the device_map="auto" mechanism is still not mature on XPU, causing model to be loaded on CPU, rather than on XPU.

If there is no particular reason for using auto, I would suggest using torch_device instead auto, because torch_device is more specific than auto and we don't have the need to use auto, e.g. large model inference, in our tests anyway. WDYT? @ArthurZucker

[amyeroberts]

Hi @faaany - thanks for opening this PR! Using device_map="auto" is necessary for this test - it's checking that beam search works when the model is split across devices. If it doesn't work with XPU, then you can add a skip with unittest e.g.

if "xpu" in torch_device:
    return unittest.skip("device_map='auto' does not work with XPU devices")

[faaany]

Hi @amyeroberts , thanks so much for reviewing this PR! I updated my patch and used torch_device=="xpu" instead of 'xpu' in torch_device, hope this is fine for you. And yes, I will remove this skip once auto works on XPU.

[amyeroberts]

Thanks for handling!

Just a small nit

[faaany]

Thanks for handling!

Just a small nit

done, thx!

[amyeroberts]

Thanks!