[tests] fix bnb test bug

[faaany]

What does this PR do?

Currently, the below test fails:

pytest -rA tests/models/esm/test_modeling_esm.py::EsmModelIntegrationTest::test_inference_bitsandbytes

>       return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
E       RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper_CUDA__index_select)

../../miniforge3/envs/hf-dev/lib/python3.10/site-packages/torch/nn/functional.py:2551: RuntimeError
----------------------------------------------------------- Captured stderr call -----------------------------------------------------------
The `load_in_4bit` and `load_in_8bit` arguments are deprecated and will be removed in the future versions. Please, pass a `BitsAndBytesConfig` object in `quantization_config` argument instead.
`low_cpu_mem_usage` was None, now default to True since model is quantized.
Loading checkpoint shards: 100%|██████████| 2/2 [00:01<00:00,  1.17it/s]
========================================================= short test summary info ==========================================================
FAILED tests/models/esm/test_modeling_esm.py::EsmModelIntegrationTest::test_inference_bitsandbytes - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument fo...

This PR fixes it.

cc @ydshieh

[ydshieh]

You are the device expert! Thank you a lot, confirmed passing!