Compare device name of profiler dynamically (#150396)

elpis-furiosa · pytorchmergebot · commit 0da8127f77f9 · 2025-04-02T06:06:06.000Z
Compare self.use_device of torch.autograd.profiler.profiler with _get_privateuse1_backend_name(), since privateuse1 backend can be renamed. Pull Request resolved: #150396 Approved by: https://github.com/sraikund16
diff --git a/torch/autograd/profiler.py b/torch/autograd/profiler.py
@@ -629,7 +629,7 @@ def _device_memory_usage(mem_record):
             )
             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)

Original file line number	Diff line number	Diff line change
`@@ -629,7 +629,7 @@ def _device_memory_usage(mem_record):`
`629`	`629`	`)`
`630`	`630`	`max_evt_id = max(max_evt_id, fe.id)`
`631`	`631`	`if fe.device_type == DeviceType.CPU and not fe.is_async:`
`632`		`- if self.use_device == "privateuseone":`
	`632`	`+ if self.use_device == _get_privateuse1_backend_name():`
`633`	`633`	`privateuse1_time = kineto_event.privateuse1_elapsed_us()`
`634`	`634`	`if privateuse1_time > 0:`
`635`	`635`	`fe.append_kernel(fe.name, fe.device_index, privateuse1_time)`