Update nn_tutorial.py (#2368)

* Update nn_tutorial.py Fix to #1303 "add pyplot.show() in beginner tutorial." Comments to issue suggested manually commenting out pyplot.show for users not using colab. --------- Co-authored-by: Svetlana Karslioglu <svekars@fb.com>
pytorch · May 31, 2023 · 7aff96c · 7aff96c
1 parent d55a262
commit 7aff96c
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/beginner_source/nn_tutorial.py b/beginner_source/nn_tutorial.py
@@ -75,6 +75,11 @@
 import numpy as np
 
 pyplot.imshow(x_train[0].reshape((28, 28)), cmap="gray")
+# ``pyplot.show()`` only if not on Colab
+try:
+    import google.colab
+except ImportError:
+    pyplot.show()
 print(x_train.shape)
 
 ###############################################################################