Unverified Commit 876312f0 authored by Nick Hill's avatar Nick Hill Committed by GitHub
Browse files

[Core] Fix `gpu_worker.py` pre-commit errors (#35312)


Signed-off-by: default avatarNick Hill <nickhill123@gmail.com>
parent 5de98abc
......@@ -744,7 +744,8 @@ class Worker(WorkerBase):
# Create the profiler wrapper only on the first start call
if self.profiler is None:
if self.profiler_config.profiler == "torch":
profiler_type = self.profiler_config.profiler
if profiler_type == "torch":
self.profiler = TorchProfilerWrapper(
self.profiler_config,
worker_name=trace_name,
......@@ -754,9 +755,12 @@ class Worker(WorkerBase):
logger.debug(
"Starting torch profiler with trace name: %s", trace_name
)
elif self.profiler_config.profiler == "cuda":
elif profiler_type == "cuda":
self.profiler = CudaProfilerWrapper(self.profiler_config)
logger.debug("Starting CUDA profiler")
else:
logger.warning("Unrecognized profiler: %s", profiler_type)
return
self.profiler.start()
else:
# Profiler already initialized. Restart profiling but keep
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment