Update gather_p pywrapper.

0x45f · Apr 18, 2022 · b2419b6 · b2419b6
1 parent 21565b1
commit b2419b6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/paddle/autograd/primops.py b/python/paddle/autograd/primops.py
@@ -263,7 +263,7 @@ def slice_assign(x, y, axis, starts, ends, strides, out=None):
     return out
 
 
-@REGISTER_FN('gather_p', 'X', 'Y')
+@REGISTER_FN('gather_p', 'X', 'IndexTensor', 'Y')
 def gather(x, indextensor, axis, out=None):
     attrs = {'axis': axis}
     helper = LayerHelper('gather_p', **locals())