def _replace_values_by_bool_mask(data: Var, mask: Var, new_value: Union[int, float]):
        """Replace the position in data where mask has True element to new_value."""
        indices = mb.non_zero(x=mb.cast(x=mask, dtype="int32"))

        # Expand the replacement value to the compatible shape for scatter_nd.
        replacement_values = mb.expand_dims(x=new_value, axes=[0])
        reps = mb.expand_dims(x=_utils.pymil_value_at(mb.shape(x=indices), 0), axes=[0])
        replacement_values = mb.tile(x=replacement_values, reps=reps)

        # Replace all nan to the corresponding values.
        return mb.scatter_nd(data=data, indices=indices, updates=replacement_values, mode="update")
Replace the position in data where mask has True element to new_value.
coremltools/converters/mil/frontend/torch/ops.py