Minor fix required to handle a pandas bug, Pandas issue: pandas-dev/pandas#36948 and PR: pandas-dev/pandas#36950.

Also makes the set_axis call before plotting the gradient series deliberately in place.

PiperOrigin-RevId: 335933611

Author: sipple