Loss not getting printed in when manual optimization is turned on

params: <h3><a name="bug-description-1" class="anchor" href="#bug-description-1"></a>Bug description</h3>
<p>I ported an simple toy example in pytorch lightning and found that when manual optimization is turned on (to implement a custom optimizer), the loss doesn’t get printed. Not sure if I am missing to set some something.</p>
<p>Any help is appreciated.</p>
<h3><a name="what-version-are-you-seeing-the-problem-on-2" class="anchor" href="#what-version-are-you-seeing-the-problem-on-2"></a>What version are you seeing the problem on?</h3>
<p>v2.1</p>
<h3><a name="how-to-reproduce-the-bug-3" class="anchor" href="#how-to-reproduce-the-bug-3"></a>How to reproduce the bug</h3>
<p>Baseline toy example (works fine and prints loss correctly):</p>
<pre data-code-wrap="python"><code class="lang-python">
import pytorch_lightning as pl
import numpy as np
import torch
from torch.nn import MSELoss
from torch.optim import Adam
from ptadamw import AdamW
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn


class SimpleDataset(Dataset):
    def __init__(self):
        X = np.arange(10000)
        y = X * 2
        X = [[_] for _ in X]
        y = [[_] for _ in y]
        self.X = torch.Tensor(X)
        self.y = torch.Tensor(y)

    def __len__(self):
        return len(self.y)

    def __getitem__(self, idx):
        return {"X": self.X[idx], "y": self.y[idx]}


class MyModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(1, 1)
        self.criterion = MSELoss()

    def forward(self, inputs_id, labels=None):
        outputs = self.fc(inputs_id)
        loss = 0
        if labels is not None:
            loss = self.criterion(outputs, labels)
        return loss, outputs

    def train_dataloader(self):
        dataset = SimpleDataset()
        return DataLoader(dataset, batch_size=1000)

    def training_step(self, batch, batch_idx):
        input_ids = batch["X"]
        labels = batch["y"]
        loss, outputs = self(input_ids, labels)
        return {"loss": loss}

    def configure_optimizers(self):
        optimizer = Adam(self.parameters())
        return optimizer


if __name__ == '__main__':
    model = MyModel()
    trainer = pl.Trainer(max_epochs=100, gpus=1)
    trainer.fit(model)

    X = torch.Tensor([[1.0], [51.0], [89.0]])
    _, y = model(X)
    print(y)

</code></pre>
<p>Ported toy example with manual optimization (runs without any errors but no loss being printed in the progress bar):</p>
<pre><code class="lang-auto">import pytorch_lightning as pl
import numpy as np
import torch
from torch.nn import MSELoss
from torch.optim import Adam
from sam import SAM
from ptadamw import AdamW
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn


class SimpleDataset(Dataset):
    def __init__(self):
        X = np.arange(10000)
        y = X * 2
        X = [[_] for _ in X]
        y = [[_] for _ in y]
        self.X = torch.Tensor(X)
        self.y = torch.Tensor(y)

    def __len__(self):
        return len(self.y)

    def __getitem__(self, idx):
        return {"X": self.X[idx], "y": self.y[idx]}


class MyModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(1, 1)
        self.criterion = MSELoss()
        self.automatic_optimization = False

    def forward(self, inputs_id, labels=None):
        outputs = self.fc(inputs_id)
        loss = 0
        if labels is not None:
            loss = self.criterion(outputs, labels)
        return loss, outputs

    def train_dataloader(self):
        dataset = SimpleDataset()
        return DataLoader(dataset, batch_size=1000)

    def training_step(self, batch, batch_idx):
        optimizer = self.optimizers()

        # first forward-backward pass
        loss_1 = self.compute_loss(batch)
        self.manual_backward(loss_1)
        optimizer.first_step(zero_grad=True)

        # second forward-backward pass
        loss_2 = self.compute_loss(batch)
        self.manual_backward(loss_2)
        optimizer.second_step(zero_grad=True)

        return {"loss": loss_1}

    def compute_loss(self, batch):
        input_ids = batch["X"]
        labels = batch["y"]
        loss, outputs = self(input_ids, labels)
        return loss

    def configure_optimizers(self):
        base_optimizer = Adam
        optimizer = SAM(model.parameters(), base_optimizer, lr=0.01, rho=0.05)
        return optimizer


if __name__ == '__main__':
    model = MyModel()
    trainer = pl.Trainer(max_epochs=100, gpus=1)
    trainer.fit(model)

    X = torch.Tensor([[1.0], [51.0], [89.0]])
    _, y = model(X)
    print(y)

</code></pre>
<pre><code class="lang-auto">

### Error messages and logs

Baseline code: time, it/s, loss and v_num fields are printed correctly
</code></pre>
<p>Epoch 99: 100%|█████████████████████████████████████████████████████████████████| 10/10 [00:00&lt;00:00, 203.64it/s, loss=4.16e+07, v_num=32]<br>
<code>Trainer.fit</code> stopped: <code>max_epochs=100</code> reached.<br>
Epoch 99: 100%|█████████████████████████████████████████████████████████████████| 10/10 [00:00&lt;00:00, 197.27it/s, loss=4.16e+07, v_num=32]<br>
tensor([[ 1.9360],<br>
[46.3258],<br>
[80.0620]], grad_fn=)</p>
<pre><code class="lang-auto">
Ported code: (loss is not being printed)
</code></pre>
<p>Epoch 99: 100%|████████████████████████████████████████████████████████████████████████████████| 10/10 [00:00&lt;00:00, 180.96it/s, v_num=33]<br>
tensor([[  2.8524],<br>
[102.8746],<br>
[178.8915]], grad_fn=)</p>
<pre><code class="lang-auto">

### Environment
</code></pre>
<p>pytorch-lightning==1.9.5<br>
lightning==2.1.3<br>
python=3.10.13=hd12c33a_1_cpython<br>
torch==2.1.2+cu118<br>
torchaudio==2.1.2+cu118<br>
torchmetrics==1.3.0<br>
torchvision==0.16.2+cu118</p>
<pre><code class="lang-auto"></code></pre>

ptrblck: <p>You might want to cross-post the question into the lightning discussion board as it seems to be related to it’s usage.</p>