fix plotting with > 1-d tensors in tutorials with py3.9 (#2013)

Summary:
Pull Request resolved: #2013

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

Reviewed By: Balandat

Differential Revision: D49382057

fbshipit-source-id: 04ba98192b26117c762b4188fae454dca3f8899f

Author: sdaulton