Check that alteroplayout is definedbefore using it

apache · Sep 29, 2020 · 74cf118 · 74cf118
1 parent 69fe439
commit 74cf118
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/relay/transforms/transform_layout.h b/src/relay/transforms/transform_layout.h
@@ -269,9 +269,10 @@ Expr LayoutRewriter(const Call& ref_call, const Array<Expr>& new_args, const Obj
 
   // If there is no FInferCorrectLayout for the type, then we just assume the layout is correct.
   static auto finfer_layout = Op::GetAttrMap<FInferCorrectLayout>("FInferCorrectLayout");
+  static auto falter_layout = Op::GetAttrMap<FTVMAlterOpLayout>("FTVMAlterOpLayout");
   if (ref_call->op.as<OpNode>()) {
     Op op = Downcast<Op>(ref_call->op);
-    if (!finfer_layout.count(op)) {
+    if (falter_layout.count(op) && !finfer_layout.count(op)) {
       return memorizer.CallWithNewLayouts(ref_call, normal_new_args);
     }
   }