[Hardware] Replace torch.cuda.device_count/current_device/set_device API (#36145)
Signed-off-by: Kunshang Ji <jikunshang95@gmail.com> Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
This commit is contained in:
@@ -141,7 +141,7 @@ def get_attention_backend_params() -> list[str]:
|
||||
|
||||
|
||||
def get_tp_size_params() -> list[pytest.param]:
|
||||
num_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 1
|
||||
num_gpus = torch.accelerator.device_count() if torch.cuda.is_available() else 1
|
||||
return [pytest.param(tp, id=f"tp{tp}") for tp in TP_SIZES if tp <= num_gpus]
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user