[tests] use torch_device instead of 0 for device check

[faaany]

What does this PR do?

Below is the rest result for tests/test_big_modeling.py::BigModelingTester::test_cpu_offload_with_hook on XPU:

===================================================== test session starts =====================================================
platform linux -- Python 3.9.0, pytest-8.0.0, pluggy-1.5.0
rootdir: /home/fanli/workspace/accelerate
plugins: subtests-0.12.1, excel-1.6.0, env-1.1.3, xdist-3.6.1
collected 1 item                                                                                                              

tests/test_big_modeling.py F                                                                                            [100%]

========================================================== FAILURES ===========================================================
________________________________________ BigModelingTester.test_cpu_offload_with_hook _________________________________________

self = <test_big_modeling.BigModelingTester testMethod=test_cpu_offload_with_hook>

    @require_non_cpu
    def test_cpu_offload_with_hook(self):
        model1 = torch.nn.Linear(4, 5)
        model1, hook1 = cpu_offload_with_hook(model1)
        assert model1.weight.device == torch.device("cpu")
    
        inputs = torch.randn(3, 4)
        outputs = model1(inputs)
>       assert outputs.device == torch.device(0)
E       AssertionError: assert device(type='xpu', index=0) == device(type='cuda', index=0)
E        +  where device(type='xpu', index=0) = tensor([[ 0.4093, -0.1815,  0.2612,  1.4389,  0.2595],\n        [ 1.0791,  0.4252, -0.6515,  1.2890,  0.9664],\n        [-0.2314, -0.8029,  0.6100,  0.4829, -0.0271]], device='xpu:0',\n       grad_fn=<AddmmBackward0>).device
E        +  and   device(type='cuda', index=0) = <class 'torch.device'>(0)
E        +    where <class 'torch.device'> = torch.device

tests/test_big_modeling.py:871: AssertionError

Not sure why these tests got passed on NPU (PR #2602), but I think the fix in this PR is the right one.

@SunMarc and @muellerzr

[SunMarc]

Thanks for the PR @faaany !

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[faaany]

Thx for the good suggestion! Code updated, could you help retrigger the CI? @SunMarc

[muellerzr]

Thanks! LG2M as well :)