From 93446be0f1a678d5894c6d6b899fa10abfcafe33 Mon Sep 17 00:00:00 2001 From: Aurelius84 Date: Tue, 8 Jun 2021 19:34:13 +0800 Subject: [PATCH] [Dy2Stat]move data to CUDAPlace in advance (#33345) * move data to CUDAPlace in advance --- .../dygraph/dygraph_to_static/partial_program.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/python/paddle/fluid/dygraph/dygraph_to_static/partial_program.py b/python/paddle/fluid/dygraph/dygraph_to_static/partial_program.py index 6eea883226b..719b06c659f 100644 --- a/python/paddle/fluid/dygraph/dygraph_to_static/partial_program.py +++ b/python/paddle/fluid/dygraph/dygraph_to_static/partial_program.py @@ -16,6 +16,7 @@ from __future__ import print_function import numpy as np import six +import paddle from paddle.fluid import framework, backward, core from paddle.fluid.dygraph import layers from paddle.fluid.dygraph.base import switch_to_static_graph @@ -263,7 +264,17 @@ class PartialProgramLayer(layers.Layer): place=framework._current_expected_place(), zero_copy=True) elif isinstance(value, core.VarBase): - var = value + if value.stop_gradient: + # NOTE(Aurelius84): If var is on CPUPlace, it will be transformed multi times + # into CUDAPlace when it's as input of multi Ops. so we move it in advance + # to avoid this problem. + var = paddle.to_tensor( + value, + dtype=value.dtype, + place=framework._current_expected_place(), + stop_gradient=True) + else: + var = value var.name = self._inputs[i].desc.name() else: continue -- GitLab