[checkpoint] rebase

apache · Sep 24, 2021 · f64e372 · f64e372
1 parent ebc116e
commit f64e372
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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) {