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: e75494aab138e1a7ea67d5bcfab8a8698f77ab6c

Author: sdaulton