deploy: SimplexLab/TorchJD@e760f22

Author: ValerianRey

latest/_sources/examples/lightning_integration.rst.txt

@@ -1,13 +1,13 @@
 PyTorch Lightning Integration
 =============================
 
-To use Jacobian descent with TorchJD in a :class:`~lightning.LightningModule`, you need to turn off
-automatic optimization by setting ``automatic_optimization`` to ``False`` and to customize the
-``training_step`` method to make it call the appropriate TorchJD method (:doc:`backward
-<../docs/autojac/backward>` or :doc:`mtl_backward <../docs/autojac/mtl_backward>`).
+To use Jacobian descent with TorchJD in a :class:`~lightning.pytorch.core.LightningModule`, you need
+to turn off automatic optimization by setting ``automatic_optimization`` to ``False`` and to
+customize the ``training_step`` method to make it call the appropriate TorchJD method
+(:doc:`backward <../docs/autojac/backward>` or :doc:`mtl_backward <../docs/autojac/mtl_backward>`).
 
 The following code example demonstrates a basic multi-task learning setup using a
-:class:`~lightning.LightningModule` that will call :doc:`mtl_backward
+:class:`~lightning.pytorch.core.LightningModule` that will call :doc:`mtl_backward
 <../docs/autojac/mtl_backward>` at each training iteration.
 
 .. code-block:: python

latest/examples/lightning_integration/index.html

@@ -240,11 +240,12 @@
         <article role="main">
           <section id="pytorch-lightning-integration">
 <h1>PyTorch Lightning Integration<a class="headerlink" href="#pytorch-lightning-integration" title="Link to this heading">¶</a></h1>
-<p>To use Jacobian descent with TorchJD in a <code class="xref py py-class docutils literal notranslate"><span class="pre">LightningModule</span></code>, you need to turn off
-automatic optimization by setting <code class="docutils literal notranslate"><span class="pre">automatic_optimization</span></code> to <code class="docutils literal notranslate"><span class="pre">False</span></code> and to customize the
-<code class="docutils literal notranslate"><span class="pre">training_step</span></code> method to make it call the appropriate TorchJD method (<a class="reference internal" href="../../docs/autojac/backward/"><span class="doc">backward</span></a> or <a class="reference internal" href="../../docs/autojac/mtl_backward/"><span class="doc">mtl_backward</span></a>).</p>
+<p>To use Jacobian descent with TorchJD in a <a class="reference external" href="https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.core.LightningModule.html#lightning.pytorch.core.LightningModule" title="(in PyTorch Lightning v2.5.1.post0)"><code class="xref py py-class docutils literal notranslate"><span class="pre">LightningModule</span></code></a>, you need
+to turn off automatic optimization by setting <code class="docutils literal notranslate"><span class="pre">automatic_optimization</span></code> to <code class="docutils literal notranslate"><span class="pre">False</span></code> and to
+customize the <code class="docutils literal notranslate"><span class="pre">training_step</span></code> method to make it call the appropriate TorchJD method
+(<a class="reference internal" href="../../docs/autojac/backward/"><span class="doc">backward</span></a> or <a class="reference internal" href="../../docs/autojac/mtl_backward/"><span class="doc">mtl_backward</span></a>).</p>
 <p>The following code example demonstrates a basic multi-task learning setup using a
-<code class="xref py py-class docutils literal notranslate"><span class="pre">LightningModule</span></code> that will call <a class="reference internal" href="../../docs/autojac/mtl_backward/"><span class="doc">mtl_backward</span></a> at each training iteration.</p>
+<a class="reference external" href="https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.core.LightningModule.html#lightning.pytorch.core.LightningModule" title="(in PyTorch Lightning v2.5.1.post0)"><code class="xref py py-class docutils literal notranslate"><span class="pre">LightningModule</span></code></a> that will call <a class="reference internal" href="../../docs/autojac/mtl_backward/"><span class="doc">mtl_backward</span></a> at each training iteration.</p>
 <div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span><span class="w"> </span><span class="nn">torch</span>
 <span class="kn">from</span><span class="w"> </span><span class="nn">lightning</span><span class="w"> </span><span class="kn">import</span> <span class="n">LightningModule</span><span class="p">,</span> <span class="n">Trainer</span>
 <span class="kn">from</span><span class="w"> </span><span class="nn">lightning.pytorch.utilities.types</span><span class="w"> </span><span class="kn">import</span> <span class="n">OptimizerLRScheduler</span>