deploy: SimplexLab/TorchJD@1684cbd

Author: ValerianRey

latest/examples/lightning_integration/index.html

@@ -247,12 +247,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 <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.4)"><code class="xref py py-class docutils literal notranslate"><span class="pre">LightningModule</span></code></a>, you need
+<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.5)"><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
-<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.4)"><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>
+<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.5)"><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>