data_parallel_tutorial to run on multigpu (#2477)

clee2000 · malfet · web-flow · commit dcb527d0432c · 2023-06-27T14:57:45.000-07:00
* data_parallel_tutorial to run on multigpu

* Collect unused memory

* set numpy seed

* empty cuda cache

* update

---------

Co-authored-by: Nikita Shulga &lt;nshulga@meta.com&gt;
diff --git a/.jenkins/metadata.json b/.jenkins/metadata.json
@@ -22,6 +22,9 @@
     "duration": 320,
     "needs": "gpu.nvidia.small.multi"
   },
+  "beginner_source/blitz/data_parallel_tutorial.py": {
+    "needs": "gpu.nvidia.small.multi"
+  },
   "intermediate_source/model_parallel_tutorial.py": {
     "needs": "gpu.nvidia.small.multi"
   },
diff --git a/conf.py b/conf.py
@@ -33,6 +33,8 @@
 sys.path.insert(0, os.path.abspath('./.jenkins'))
 import pytorch_sphinx_theme
 import torch
+import numpy
+import gc
 import glob
 import random
 import shutil
@@ -91,6 +93,8 @@ def reset_seeds(gallery_conf, fname):
     torch.manual_seed(42)
     torch.set_default_device(None)
     random.seed(10)
+    numpy.random.seed(10)
+    gc.collect()
 
 sphinx_gallery_conf = {
     'examples_dirs': ['beginner_source', 'intermediate_source',
