diff --git a/beginner_source/introyt/autogradyt_tutorial.py b/beginner_source/introyt/autogradyt_tutorial.py
index abf75a7d266..c065aaebd56 100644
--- a/beginner_source/introyt/autogradyt_tutorial.py
+++ b/beginner_source/introyt/autogradyt_tutorial.py
@@ -473,16 +473,14 @@ def add_tensors2(x, y):
 # 
 
 device = torch.device('cpu')
-run_on_gpu = False
 if torch.cuda.is_available():
     device = torch.device('cuda')
-    run_on_gpu = True
-    
+
 x = torch.randn(2, 3, requires_grad=True)
 y = torch.rand(2, 3, requires_grad=True)
 z = torch.ones(2, 3, requires_grad=True)
 
-with torch.autograd.profiler.profile(use_cuda=run_on_gpu) as prf:
+with torch.autograd.profiler.profile(use_device=device.type) as prf:
     for _ in range(1000):
         z = (z / x) * y