From a31cf3d2764e310759cb545bb39619dce7a13530 Mon Sep 17 00:00:00 2001
From: Leo Chen <chenqiuliang@baidu.com>
Date: Thu, 31 Aug 2023 11:45:19 +0800
Subject: [PATCH] skip data_transfer for save op (#56775)

---
 paddle/fluid/operators/save_op.cc | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/paddle/fluid/operators/save_op.cc b/paddle/fluid/operators/save_op.cc
index bc1f5a0d34f..2450d64542a 100644
--- a/paddle/fluid/operators/save_op.cc
+++ b/paddle/fluid/operators/save_op.cc
@@ -101,7 +101,9 @@ PD_REGISTER_KERNEL(save,
                    int16_t,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+}
 
 PD_REGISTER_KERNEL(save_sr,
                    CPU,
@@ -115,7 +117,9 @@ PD_REGISTER_KERNEL(save_sr,
                    int16_t,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+}
 
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
 PD_REGISTER_KERNEL(save,
@@ -130,7 +134,9 @@ PD_REGISTER_KERNEL(save,
                    int16_t,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+}
 
 PD_REGISTER_KERNEL(save_sr,
                    GPU,
@@ -144,5 +150,7 @@ PD_REGISTER_KERNEL(save_sr,
                    int16_t,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+}
 #endif
-- 
GitLab