fix plotting with > 1-d tensors in tutorials with py3.9

Summary: see title. This is causing a failure in the tutorials on  #2011

Differential Revision: D49382057

fbshipit-source-id: f1c62d17d7485ce7e2a381646fad84885ce0d94b

Author: sdaulton