update closure description

Author: awaelchli

docs/source/common/optimizers.rst

@@ -299,10 +299,14 @@ Under the hood, Lightning does the following:
 
     for epoch in epochs:
         for batch in data:
-            loss = model.training_step(batch, batch_idx, ...)
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
+
+            def closure():
+                loss = model.training_step(batch, batch_idx, ...)
+                optimizer.zero_grad()
+                loss.backward()
+                return loss
+
+            optimizer.step(closure)
 
         for lr_scheduler in lr_schedulers:
             lr_scheduler.step()
@@ -314,14 +318,22 @@ In the case of multiple optimizers, Lightning does the following:
     for epoch in epochs:
         for batch in data:
             for opt in optimizers:
-                loss = model.training_step(batch, batch_idx, optimizer_idx)
-                opt.zero_grad()
-                loss.backward()
-                opt.step()
+
+                def closure():
+                    loss = model.training_step(batch, batch_idx, optimizer_idx)
+                    opt.zero_grad()
+                    loss.backward()
+                    return loss
+
+                opt.step(closure)
 
         for lr_scheduler in lr_schedulers:
             lr_scheduler.step()
 
+As can be seen in the code snippet above, Lightning defines a closure with ``training_step``, ``zero_grad``
+and ``backward`` for the optimizer to execute. This mechanism is in place to support optimizers which operate on the
+output of the closure (e.g. the loss) or need to call the closure several times (e.g. :class:`~torch.optim.LBFGS`).
+
 .. warning::
    Before 1.2.2, Lightning internally calls ``backward``, ``step`` and ``zero_grad`` in the order.
    From 1.2.2, the order is changed to ``zero_grad``, ``backward`` and ``step``.