diff --git a/paddle/operators/get_places_op.cc b/paddle/operators/get_places_op.cc index dd937488f4cb2b8ac53a6b0b179416ab1faedee5..96a019ac799d3477f7319562ab0f7bc26f8a6560 100644 --- a/paddle/operators/get_places_op.cc +++ b/paddle/operators/get_places_op.cc @@ -35,7 +35,7 @@ class GetPlacesOp : public framework::OperatorBase { out_var_name); auto &places = *(out_var->GetMutable>()); - places.reserve(trainer_count); + places.resize(trainer_count); if (use_gpu) { for (int i = 0; i < trainer_count; i++) { places.emplace_back(platform::GPUPlace(i)); diff --git a/python/paddle/v2/fluid/tests/test_layers.py b/python/paddle/v2/fluid/tests/test_layers.py index 2286e94a90a4810dfb170ba6e929a7c4f3edaba1..4a03ca68fffca556ba3747f64bb858d0e735d33e 100644 --- a/python/paddle/v2/fluid/tests/test_layers.py +++ b/python/paddle/v2/fluid/tests/test_layers.py @@ -170,6 +170,12 @@ class TestBook(unittest.TestCase): self.assertIsNotNone(layers.sequence_expand(x=x, y=y)) print(str(program)) + def test_get_places(self): + program = Program() + with program_guard(program): + x = layers.get_places(use_gpu=True, trainer_count=4) + print(str(program)) + if __name__ == '__main__': unittest.main()