7
7
"""
8
8
9
9
######################################################################
10
- # Summary
11
- # ~~~~~~~~
12
10
#
13
- # In this tutorial we will apply torch.compile to the optimizer to observe
14
- # the GPU performance improvement
11
+ # In this recipe, we will apply `` torch.compile`` to the optimizer to observe
12
+ # the GPU performance improvement.
15
13
#
16
14
# .. note::
17
15
#
22
20
######################################################################
23
21
# Model Setup
24
22
# ~~~~~~~~~~~~~~~~~~~~~
25
- # For this example we'll use a simple sequence of linear layers.
26
- # Since we are only benchmarking the optimizer, choice of model doesn't matter
23
+ # For this example, we'll use a simple sequence of linear layers.
24
+ # Since we are only benchmarking the optimizer, the choice of model doesn't matter
27
25
# because optimizer performance is a function of the number of parameters.
28
26
#
29
27
# Depending on what machine you are using, your exact results may vary.
39
37
40
38
#############################################################################
41
39
# Setting up and running the optimizer benchmark
42
- # ~~~~~~~~~~~~~~~~~~~~~~~~~
40
+ # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
43
41
# In this example, we'll use the Adam optimizer
44
42
# and create a helper function to wrap the step()
45
43
# in torch.compile()
@@ -52,7 +50,7 @@ def fn():
52
50
opt .step ()
53
51
54
52
55
- # Lets define a helpful benchmarking function:
53
+ # Let's define a helpful benchmarking function:
56
54
import torch .utils .benchmark as benchmark
57
55
58
56
0 commit comments