Fix param@grad type error for amp in run_program (PaddlePaddle#40938)

Author: 0x45f

python/paddle/fluid/dygraph/dygraph_to_static/partial_program.py

@@ -204,7 +204,9 @@ def _train_amp_program(self):
  """
  Lazy initialized property of train_amp_program.
  """
- return self._append_backward_desc(self._infer_amp_program)
+ train_amp_program = self._append_backward_desc(self._infer_amp_program)
+ self._set_grad_type(self._params, train_amp_program)
+ return train_amp_program
 
  @LazyInitialized
  @switch_to_static_graph
@@ -224,7 +226,10 @@ def _train_pure_fp16_program(self):
  """
  Lazy initialized property of _train_pure_fp16_program.
  """
- return self._append_backward_desc(self._infer_pure_fp16_program)
+ train_pure_fp16_program = self._append_backward_desc(
+ self._infer_pure_fp16_program)
+ self._set_grad_type(self._params, train_pure_fp16_program)
+ return train_pure_fp16_program
 
  @LazyInitialized
  def _infer_program_id(self):