If you are just using it to remove the burn-in, you can do pm.traceplot(trace[burn_in:]) instead.
pm.traceplot(trace[burn_in:])