Fix onnx output for GcnTransform (#4786)

dotnet · Feb 5, 2020 · 46e7dc6 · 46e7dc6
1 parent d2ff3e3
commit 46e7dc6
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/Microsoft.ML.Transforms/GcnTransform.cs b/src/Microsoft.ML.Transforms/GcnTransform.cs
@@ -617,7 +617,8 @@ public void SaveAsOnnx(OnnxContext ctx)
                         continue;
                     }
 
-                    if (!SaveAsOnnxCore(ctx, iinfo, ctx.GetVariableName(inputColumnName), ctx.AddIntermediateVariable(_srcTypes[iinfo], inputColumnName)))
+                    string outputColumnName = _parent.ColumnPairs[iinfo].outputColumnName;
+                    if (!SaveAsOnnxCore(ctx, iinfo, ctx.GetVariableName(inputColumnName), ctx.AddIntermediateVariable(_srcTypes[iinfo], outputColumnName)))
                     {
                         ctx.RemoveColumn(inputColumnName, true);
                     }