diff --git a/python/paddle/fluid/dygraph/dygraph_to_static/utils.py b/python/paddle/fluid/dygraph/dygraph_to_static/utils.py index d299e63fd0073ed5048bbd5ba038c92db6f9079d..3f42137791710edd5b61ce3798ec79838210b3ee 100644 --- a/python/paddle/fluid/dygraph/dygraph_to_static/utils.py +++ b/python/paddle/fluid/dygraph/dygraph_to_static/utils.py @@ -882,6 +882,8 @@ class ForNodeVisitor(object): self.node.iter.func, gast.Attribute) and self.node.iter.func.attr == 'numpy': return True + elif isinstance(self.node.iter, gast.Subscript): + return True else: return False diff --git a/python/paddle/fluid/tests/unittests/dygraph_to_static/test_for_enumerate.py b/python/paddle/fluid/tests/unittests/dygraph_to_static/test_for_enumerate.py index a74c56fc31766ccf39cbcbee6b1138573fe9de6a..18995238a3c0589328bcc838f1cb7a04acbf61d7 100644 --- a/python/paddle/fluid/tests/unittests/dygraph_to_static/test_for_enumerate.py +++ b/python/paddle/fluid/tests/unittests/dygraph_to_static/test_for_enumerate.py @@ -159,6 +159,7 @@ def for_enumerate_var_numpy_with_start_continue(x_array): def for_iter_var(x_array): z = fluid.layers.fill_constant([1], 'int32', 0) x_array = fluid.dygraph.to_variable(x_array) + for x in x_array: z = z + x return z @@ -221,6 +222,17 @@ def for_enumerate_var_with_nested_range(x_array): return x +# 16. for iter var[idx] +@paddle.jit.to_static +def for_iter_var_idx(x_array): + z = fluid.layers.fill_constant([1], 'int32', 0) + x_array = fluid.dygraph.to_variable(x_array) + + for x in x_array[0:]: + z = z + x + return z + + class TestTransformBase(unittest.TestCase): def setUp(self): self.place = fluid.CUDAPlace(0) if fluid.is_compiled_with_cuda( @@ -343,6 +355,11 @@ class TestForIterVar(TestForIterVarNumpy): self.dygraph_func = for_iter_var +class TestForIterVarIdx(TestForIterVarNumpy): + def set_test_func(self): + self.dygraph_func = for_iter_var_idx + + class TestForEnumerateVar(TestForIterVarNumpy): def set_test_func(self): self.dygraph_func = for_enumerate_var