Box/Lasso selection giving wrong indices in point_inds for px.scatter

Hi,
I am implementing lasso/box selection in a scatterplot. I am using a callback function that returns me the indices of data selected by the user. The code works fine and returns me the correct indices when I applied a continuous or no label to the points. However, when I color the points using a discrete label I get incorrect indices or empty array.

Sample code to run is given below.

import plotly.express as px
import plotly.graph_objects as go
from plotly.graph_objs import *
from IPython.display import display

df = px.data.iris()
selected = []
def filter_(trace, points, state):
    global selected
    print(points)
    selected = points.point_inds

layout = Layout(
    paper_bgcolor='rgba(0,0,0,0)',
    plot_bgcolor='rgba(0,0,0,0)',
    width=1000,
    height=600,

)

fig = px.scatter(
    df,
    x='sepal_length',
    y='sepal_width',
    color="species",

)

fig.layout=layout
fig.update_xaxes(visible=False)
fig.update_yaxes(visible=False)
fig = go.FigureWidget(fig.data, fig.layout)
fig.layout.dragmode = 'lasso'
fig.data[0].on_selection(filter_)

display(fig)
dh = display(display_id=True)

On making a lasso/box selection this is the output thrown. As you can see point_inds is empty.

Points(point_inds=[],
       xs=[],
       ys=[],
       trace_name='setosa',
       trace_index=0)

If you comment color="species" , the code works fine.

In the above code, I also applied on_selection on all objects of fig.data but still did not work.

for f in fig.data:
    f.on_selection(filter_)

Version info -

Python - '3.8.8'
Plotly - '5.11.0'

Reproduce the provided px.scatter and FigureWidget example in a Jupyter notebook using Python 3.8.8 and Plotly 5.11.0, comparing selection with and without color="species". Trace the selection callbacks on fig.data and verify that point_inds contains the selected rows for each colored trace; done means lasso and box selection report the expected indices instead of an empty or incorrect array.