MAINT: allow cpu and cuda targets in e2e examples

ev-br · ev-br · commit 74b98dba3445 · 2023-05-01T12:41:31.000+03:00
diff --git a/e2e/mandelbrot/mandelbrot.png b/e2e/mandelbrot/mandelbrot.png
diff --git a/e2e/mandelbrot/mandelbrot.py b/e2e/mandelbrot/mandelbrot.py
@@ -6,8 +6,11 @@
 #import numpy as np
 import torch_np as np
 
+
+# To run on CUDA, change "cpu" to "cuda" below.
 import torch
-torch.set_default_device("cuda")
+torch.set_default_device("cpu")
+
 
 # from mandelbrot_numpy_1 import mandelbrot  # copy-paste below
 
diff --git a/e2e/maze/maze_numpy.py b/e2e/maze/maze_numpy.py
@@ -12,8 +12,10 @@
 
 _np.random.seed(1234)
 
+
+# To run on CUDA, change "cpu" to "cuda" below.
 import torch
-torch.set_default_device("cuda")
+torch.set_default_device("cpu")
 
 
 def build_maze(shape=(65,65), complexity=0.75, density = 0.50):
diff --git a/e2e/nn_from_scratch/nn.py b/e2e/nn_from_scratch/nn.py
@@ -4,8 +4,10 @@
 import numpy as _np
 import torch_np as np
 
+# To run on CUDA, change "cpu" to "cuda" below.
 import torch
-torch.set_default_device('cuda')
+torch.set_default_device("cpu")
+
 
 # Creating data set
  
diff --git a/e2e/smoke/smoke_solver.py b/e2e/smoke/smoke_solver.py
@@ -12,8 +12,9 @@
 """
 import torch_np as np
 
+# To run on CUDA, change "cpu" to "cuda" below.
 import torch
-torch.set_default_device("cuda")
+torch.set_default_device("cpu")
 
 
 def set_bnd(N, b, x):
