From a73064f2e237b9f05e4e932f0ae5fc04da9877e1 Mon Sep 17 00:00:00 2001
From: chentianyu03 <chentianyu03@baidu.com>
Date: Mon, 6 Dec 2021 18:01:10 +0800
Subject: [PATCH] pylayer support tuple/list type args (#37727)

---
 paddle/fluid/imperative/py_layer_fwd.h | 44 ++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/paddle/fluid/imperative/py_layer_fwd.h b/paddle/fluid/imperative/py_layer_fwd.h
index 1baf73ab3b9..79251d7bf7a 100644
--- a/paddle/fluid/imperative/py_layer_fwd.h
+++ b/paddle/fluid/imperative/py_layer_fwd.h
@@ -101,6 +101,28 @@ py::object PyLayerApply(const platform::Place& place, const py::handle& cls,
               "`%s` type argument can not be cast into `Tensor`.",
               ptr->ptr()->ob_type->tp_name));
         }
+      } else if (py::isinstance<py::tuple>(*ptr) ||
+                 py::isinstance<py::list>(*ptr)) {
+        try {
+          auto tuple_arg = ptr->cast<py::tuple>();
+          for (auto iter = tuple_arg.begin(); iter != tuple_arg.end(); ++iter) {
+            try {
+              auto t = iter->cast<std::shared_ptr<VarBase>>();
+              input_vars.push_back(t);
+            } catch (py::cast_error& err) {
+              PADDLE_THROW(platform::errors::InvalidArgument(
+                  "The `PyLayer.forward` function contains invalid argument, "
+                  "the "
+                  "`%s` type argument can not be cast into `Tensor`.",
+                  ptr->ptr()->ob_type->tp_name));
+            }
+          }
+        } catch (py::cast_error& err) {
+          PADDLE_THROW(platform::errors::InvalidArgument(
+              "The `PyLayer.forward` function contains invalid argument, the "
+              "`%s` type argument can not be cast into `Tensor`.",
+              ptr->ptr()->ob_type->tp_name));
+        }
       }
     }
   }
@@ -119,6 +141,28 @@ py::object PyLayerApply(const platform::Place& place, const py::handle& cls,
               "`%s` type argument can not be cast into `Tensor`.",
               ptr->second.ptr()->ob_type->tp_name));
         }
+      } else if (py::isinstance<py::tuple>(*ptr->second) ||
+                 py::isinstance<py::list>(*ptr->second)) {
+        try {
+          auto tuple_arg = ptr->second.cast<py::tuple>();
+          for (auto iter = tuple_arg.begin(); iter != tuple_arg.end(); ++iter) {
+            try {
+              auto t = iter->cast<std::shared_ptr<VarBase>>();
+              input_vars.push_back(t);
+            } catch (py::cast_error& err) {
+              PADDLE_THROW(platform::errors::InvalidArgument(
+                  "The `PyLayer.forward` function contains invalid argument, "
+                  "the "
+                  "`%s` type argument can not be cast into `Tensor`.",
+                  ptr->second.ptr()->ob_type->tp_name));
+            }
+          }
+        } catch (py::cast_error& err) {
+          PADDLE_THROW(platform::errors::InvalidArgument(
+              "The `PyLayer.forward` function contains invalid argument, the "
+              "`%s` type argument can not be cast into `Tensor`.",
+              ptr->second.ptr()->ob_type->tp_name));
+        }
       }
     }
   }
-- 
GitLab