fix bug

PaddlePaddle · Nov 14, 2023 · ba8dc2b · ba8dc2b
1 parent 773d8c2
commit ba8dc2b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py
@@ -1183,7 +1183,9 @@ def eye(num_rows, num_columns=None, dtype=None, name=None):
     """
 
     def _check_attr(attr, message):
-        if isinstance(attr, ((Variable, core.eager.Tensor))):
+        if isinstance(
+            attr, ((Variable, core.eager.Tensor, paddle.pir.OpResult))
+        ):
             assert len(attr.shape) == 1 and attr.shape[0] in [1, -1]
         elif not isinstance(attr, int) or attr < 0:
             raise TypeError(f"{message} should be a non-negative int.")