Compare device name of profiler dynamically (#150396)

Compare self.use_device of torch.autograd.profiler.profiler with _get_privateuse1_backend_name(), since privateuse1 backend can be renamed.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/150396
Approved by: https://github.com/sraikund16
This commit is contained in:
Sukchul Cho 2025-04-02 06:06:02 +00:00 committed by PyTorch MergeBot
parent c65de03196
commit 0da8127f77

View File

@ -629,7 +629,7 @@ class profile:
)
max_evt_id = max(max_evt_id, fe.id)
if fe.device_type == DeviceType.CPU and not fe.is_async:
if self.use_device == "privateuseone":
if self.use_device == _get_privateuse1_backend_name():
privateuse1_time = kineto_event.privateuse1_elapsed_us()
if privateuse1_time > 0:
fe.append_kernel(fe.name, fe.device_index, privateuse1_time)