diff --git a/python/paddle/fluid/dygraph/dygraph_to_static/utils.py b/python/paddle/fluid/dygraph/dygraph_to_static/utils.py index 4da898d7441d80..c4df01c4c7654a 100644 --- a/python/paddle/fluid/dygraph/dygraph_to_static/utils.py +++ b/python/paddle/fluid/dygraph/dygraph_to_static/utils.py @@ -1045,7 +1045,8 @@ def is_for_range_iter(self): gast.Name) and self.node.iter.func.id == "range" def is_for_iter(self): - if isinstance(self.node.iter, (gast.Name, gast.Attribute)): + if isinstance(self.node.iter, + (gast.Name, gast.Attribute, gast.List, gast.Tuple)): return True elif isinstance(self.node.iter, gast.Call) and isinstance( self.node.iter.func, 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 305bdc1468b285..2aab27c03110d1 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 @@ -304,6 +304,24 @@ def forward(self, x): return z +# 21. for original list +@paddle.jit.to_static +def for_original_list(): + z = fluid.layers.fill_constant([1], 'int32', 0) + for x in [1, 2, 3]: + z = z + x + return z + + +# 22. for original tuple +@paddle.jit.to_static +def for_original_tuple(): + z = fluid.layers.fill_constant([1], 'int32', 0) + for x in (1, 2, 3): + z = z + x + return z + + class TestTransformBase(unittest.TestCase): def setUp(self): self.place = fluid.CUDAPlace(0) if fluid.is_compiled_with_cuda( @@ -344,6 +362,13 @@ def transformed_result_compare(self): self.assertTrue(np.allclose(x.numpy(), y.numpy())) +class TestTransformForOriginalList(TestTransform): + def _run(self, to_static): + program_translator.enable(to_static) + with fluid.dygraph.guard(): + return self.dygraph_func() + + class TestTransformError(TestTransformBase): def transformed_error(self, etype): with self.assertRaises(etype): @@ -471,5 +496,21 @@ def set_test_func(self): self.dygraph_func = ForwardContainsForLayer() +class TestForOriginalList(TestTransformForOriginalList): + def set_test_func(self): + self.dygraph_func = for_original_list + + def test_transformed_result_compare(self): + self.transformed_result_compare() + + +class TestForOriginalTuple(TestTransformForOriginalList): + def set_test_func(self): + self.dygraph_func = for_original_tuple + + def test_transformed_result_compare(self): + self.transformed_result_compare() + + if __name__ == '__main__': unittest.main()