Adding clarifying documentation on the usage of second_order_closure (#3551)

* Adding clarifying documentation on the usage of second_order_closure

* oops typo

* making functions more sane

* fixing spacing issues - I think

* Apply suggestions from code review

* suggestions

Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
Co-authored-by: Rohit Gupta <rohitgr1998@gmail.com>

docs/source/optimizers.rst

@@ -123,3 +123,29 @@ Here we add a learning-rate warm up
         # update params
         optimizer.step()
         optimizer.zero_grad()
+
+----------
+
+Using the closure functions for optimization
+--------------------------------------------
+
+When using optimization schemes such as LBFGS, the `second_order_closure` needs to be enabled. By default, this function is defined by wrapping the `training_step` and the backward steps as follows
+
+.. testcode::
+
+    def second_order_closure(pl_module, split_batch, batch_idx, opt_idx, optimizer, hidden):
+        # Model training step on a given batch
+        result = pl_module.training_step(split_batch, batch_idx, opt_idx, hidden)
+
+        # Model backward pass
+        pl_module.backward(result, optimizer, opt_idx)
+
+        # on_after_backward callback
+        pl_module.on_after_backward(result.training_step_output, batch_idx, result.loss)
+
+        return result
+
+    # This default `second_order_closure` function can be enabled by passing it directly into the `optimizer.step`
+    def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_idx, second_order_closure, on_tpu=False, using_native_amp=False, using_lbfgs=False):
+        # update params
+        optimizer.step(second_order_closure)