Set output shapes after materialization.

bgogul · bgogul · commit 834d97d20722 · 2019-08-01T00:40:52.000-07:00
diff --git a/Sources/TensorFlow/Core/LazyTensorOperation.swift b/Sources/TensorFlow/Core/LazyTensorOperation.swift
@@ -891,6 +891,7 @@ extension LazyTensorOperation {
         for lazyOp in traceInfo.lazyOperations {
             let end = start + lazyOp.outputCount
             lazyOp.outputs = Array(allOutputs[start..<end])
+            lazyOp.outputShapes = lazyOp.outputs!.map { $0.shape }
             start = end
         }
 

Original file line number	Diff line number	Diff line change
`@@ -891,6 +891,7 @@ extension LazyTensorOperation {`
`891`	`891`	`for lazyOp in traceInfo.lazyOperations {`
`892`	`892`	`let end = start + lazyOp.outputCount`
`893`	`893`	`lazyOp.outputs = Array(allOutputs[start..<end])`
	`894`	`+ lazyOp.outputShapes = lazyOp.outputs!.map { $0.shape }`
`894`	`895`	`start = end`
`895`	`896`	`}`
`896`	`897`