diff --git a/conf.py b/conf.py index eaa25a956c6..5f88045adb3 100644 --- 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 index f5e50b3f0be..103560fd743 100644 --- 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