Skip to content

Commit 204044a

Browse files
James Reedezyang
James Reed
authored andcommitted
Symbolic representation for unfold using ATen (pytorch#3334)
1 parent ac8f566 commit 204044a

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

torch/onnx/symbolic.py

+4
Original file line numberDiff line numberDiff line change
@@ -233,3 +233,7 @@ def avg_pool2d(g, input, kernel_size, stride, padding, ceil_mode, count_include_
233233

234234
def log_softmax(g, input, dim=None):
235235
return g.op("Log", g.op('Softmax', input, axis_i=dim).setTypeAs(input))
236+
237+
238+
def unfold(g, input, dimension, size, step):
239+
return g.op("ATen", input, operator_s="unfold", dimension_i=dimension, size_i=size, step_i=step)

0 commit comments

Comments
 (0)