fix dynamic shape inference in DepthToSpace

Signed-off-by: Dom Miketa <[email protected]>
onnx · Mar 15, 2021 · 133db14 · 133db14
1 parent 3c358ad
commit 133db14
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/onnx_tf/handlers/backend/depth_to_space.py b/onnx_tf/handlers/backend/depth_to_space.py
@@ -40,7 +40,8 @@ def _common(cls, node, **kwargs):
     if mode == "CRD":
       # need native computation
       bsize = attrs.get("blocksize")
-      batch, channel, height, width = x.shape
+      x_shape = tf.shape(x)
+      batch, channel, height, width = x_shape[0], x_shape[1], x_shape[2], x_shape[3]
       csize = channel // (bsize**2)
 
       reshape_node = tf.reshape(x, [batch, csize, bsize, bsize, height, width])