diff options
Diffstat (limited to 'tensorflow/core/profiler/internal/tfprof_node.cc')
| -rw-r--r-- | tensorflow/core/profiler/internal/tfprof_node.cc | 4 |
1 files changed, 2 insertions, 2 deletions
diff --git a/tensorflow/core/profiler/internal/tfprof_node.cc b/tensorflow/core/profiler/internal/tfprof_node.cc index 70b91c37e4..d4a784ffaa 100644 --- a/tensorflow/core/profiler/internal/tfprof_node.cc +++ b/tensorflow/core/profiler/internal/tfprof_node.cc @@ -25,7 +25,7 @@ bool CountAsAcceleratorTime(const string& device) { } bool CountAsCPUTime(const string& device) { - return RE2::FullMatch(device, ".*/(gpu|cpu|device:sycl):\\d+"); + return RE2::FullMatch(device, ".*/(device:gpu|gpu|cpu|device:sycl):\\d+"); } bool IsCanonicalDevice(const string& device) { return CountAsCPUTime(device); } @@ -143,7 +143,7 @@ void TFGraphNode::AddStepStat(int64 step, const string& device, // TODO(xpan): Make this more robust? // See run_metadata_test.py - // It can be /job:0/replica:0/xxxx/gpu:0, or simply /gpu:0. + // It can be /job:0/replica:0/xxxx/device:GPU:0, or simply /device:GPU:0. // It can has some ad-hoc suffix, such as /stream:xx or /memcpy:xx. if (IsCanonicalDevice(dev)) { if (!canonical_device_.empty()) { |