From f64e372367ecd95a1dad6215c7b3ee2ba1f516ac Mon Sep 17 00:00:00 2001
From: Mark Shields <mbs@octoml.ai>
Date: Thu, 23 Sep 2021 14:25:43 -0700
Subject: [PATCH] [checkpoint] rebase

---
 src/relay/transforms/device_planner.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/relay/transforms/device_planner.cc b/src/relay/transforms/device_planner.cc
index 3d6c59d881258..cdbfcc78c5662 100644
--- a/src/relay/transforms/device_planner.cc
+++ b/src/relay/transforms/device_planner.cc
@@ -256,6 +256,7 @@
 #include <tvm/relay/analysis.h>
 #include <tvm/relay/attrs/annotation.h>
 #include <tvm/relay/attrs/device_copy.h>
+#include <tvm/relay/attrs/function.h>
 #include <tvm/relay/attrs/memory.h>
 #include <tvm/relay/expr_functor.h>
 #include <tvm/relay/op.h>
@@ -1113,7 +1114,7 @@ class DeviceAnalyzer : public ExprVisitor {
     // If the function already has an "on_device" attribute then we can further
     // constrain the function's domain to match it.
     Optional<Attrs> opt_attrs =
-        function_node->GetAttr<Attrs>(FunctionOnDeviceAttrs::kFunctionAttrsKey);
+        function_node->GetAttr<Attrs>(attr::kFunctionAttrsKey);
     if (opt_attrs) {
       std::vector<DeviceDomainPtr> args_and_result;
       for (size_t i = 0; i < function_node->params.size(); ++i) {