Zehao Shi
Zehao Shi

Reputation: 99

When writing a custom op, if the output shape is dynamic, how to define the output of infer_shape function?

I'm now writing a custom op, and it's output shape is dynamic, how to define the output_shape in infer_shape function.

Upvotes: 0

Views: 115

Answers (1)

Sergei
Sergei

Reputation: 1637

The infer_shape function of CustomOpProp has an in_shape parameter. You can use it to calculate the output_shape dynamically.

If you cannot infer input shape without input data, but can infer maximum shape, then you can use it as commented here: https://github.com/apache/incubator-mxnet/issues/9758

Upvotes: 1

Related Questions