fix masked_scatter (#60835)

yangguohao · web-flow · commit a92999d0788a · 2024-01-17T11:20:18.000+08:00
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -4503,12 +4503,13 @@ def masked_scatter(x, mask, value, name=None):
     zeros_like_x = paddle.zeros_like(x, dtype=int)
     mask = paddle.add(paddle.cast(mask, dtype="int"), zeros_like_x)
     mask_prefix = paddle.clip(mask.cumsum() - 1, min=0)
-    assert (
-        mask_prefix[-1] <= value.numel()
-    ), f'mask true nums must be <= value size, but got mask true nums is {mask.sum().item()}, value size is {value.numel().item()}'
+    if in_dynamic_mode():
+        assert (
+            mask_prefix[-1] <= value.numel()
+        ), f'mask true nums must be <= value size, but got mask true nums is {mask_prefix[-1].item()}, value size is {value.numel().item()}'
 
     value = value.flatten()[mask_prefix].reshape(mask.shape)
-    mask = paddle.logical_not(mask)
+    mask = paddle.logical_not(mask.astype(bool))
     return paddle.where(mask, x, value)
 
 
@@ -4530,7 +4531,7 @@ def masked_scatter_(x, mask, value, name=None):
     ), f'mask true nums must be <= value size, but got mask true nums is {mask_prefix[-1].item()}, value size is {value.numel().item()}'
 
     value = value.flatten()[mask_prefix].reshape(mask.shape)
-    mask = paddle.logical_not(mask)
+    mask = paddle.logical_not(mask.astype(bool))
     out = paddle.where_(mask, x, value)
     return out
 
