Color with categorical data

Hello,

I have a 3D plotly express figure colored with a categorical data. I made a dropdown menu to recolor the traces and toggle between lines and marker.
However, the layout update is not working if I change the color to another categorical data (lithostrat_formation).
How can I give the user multiple options to color the traces with categorical data?
the code below:

Add dropdown

fig.update_layout(
    updatemenus=[
        dict(
            buttons=list([
                dict(
                    args=["mode", "lines"],
                    label="Lines",
                    method="restyle"
                ),
                dict(
                    args=["mode", "lines+markers"],
                    label="Lines & Markers",
                    method="restyle"
                ),
                dict(
                    args=["mode", "markers"],
                    label="markers",
                    method="restyle"
                )
            ]),
            direction="down",
            pad={"r": 10, "t": 10},
            showactive=True,
            x=0.1,
            xanchor="left",
            y=1.1,
            yanchor="top"
        ),
        dict(
            buttons=list([
                dict(
                    args=[{'marker.color': [data['tvd']], 'line.color':[data['tvd']]}],
                    label="TVD",
                    method="restyle"
                ),
                dict(
                    args=[{'marker.name': [data['lithostrat_formation']], 'line.name':[data['lithostrat_formation']]}],
                    label="lithostrat_formation",
                    method="restyle"
                ),
                dict(
                    args=[{'marker.color': [data['dy']], 'line.color':[data['dy']]}],
                    label="Dy",
                    method="restyle"
                )
            ]),
            direction="down",
            pad={"r": 10, "t": 10},
            showactive=True,
            x=0.4,
            xanchor="left",
            y=1.1,
            yanchor="top"
        )
    ]
)