Set random seed (#2438)

clee2000 · malfet · web-flow · commit eaa2e901bf28 · 2023-06-07T10:27:29.000-07:00
To make tutorial builds predictable, but still keep randomness when one rans it on Collab.
Also, reset default_device after every tutorial runCo-authored-by: Nikita Shulga &lt;nshulga@meta.com&gt;

Co-authored-by: Nikita Shulga &lt;nshulga@meta.com&gt;
diff --git a/conf.py b/conf.py
@@ -34,6 +34,7 @@
 import pytorch_sphinx_theme
 import torch
 import glob
+import random
 import shutil
 from custom_directives import IncludeDirective, GalleryItemDirective, CustomGalleryItemDirective, CustomCalloutItemDirective, CustomCardItemDirective
 import distutils.file_util
@@ -85,6 +86,11 @@
 
 # -- Sphinx-gallery configuration --------------------------------------------
 
+def reset_seeds(gallery_conf, fname):
+    torch.manual_seed(42)
+    torch.set_default_device(None)
+    random.seed(10)
+
 sphinx_gallery_conf = {
     'examples_dirs': ['beginner_source', 'intermediate_source',
                       'advanced_source', 'recipes_source', 'prototype_source'],
@@ -94,7 +100,8 @@
     'backreferences_dir': None,
     'first_notebook_cell': ("# For tips on running notebooks in Google Colab, see\n"
                             "# https://pytorch.org/tutorials/beginner/colab\n"
-                            "%matplotlib inline")
+                            "%matplotlib inline"),
+    'reset_modules': (reset_seeds)
 }
 
 if os.getenv('GALLERY_PATTERN'):
diff --git a/recipes_source/recipes/changing_default_device.py b/recipes_source/recipes/changing_default_device.py
@@ -43,9 +43,6 @@
 print(mod.weight.device)
 print(mod(torch.randn(128, 20)).device)
 
-# And then globally return it back to CPU
-torch.set_default_device('cpu')
-
 ################################################################
 # This function imposes a slight performance cost on every Python
 # call to the torch API (not just factory functions). If this
