Repository: Lightning-AI/pytorch-lightning
File: tests/models/test_cpu.py (Python)

Code before:

        def __len__(self):
            return 1

    class BpttTestModel(BoringModel):
        def __init__(self, batch_size, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self.test_hidden = None
            self.batch_size = batch_size

        def training_step(self, batch, batch_idx, hiddens):
            assert hiddens == self.test_hidden, "Hidden state not persistent between tbptt steps"
            self.test_hidden = torch.rand(1)

            x_tensor, y_list = batch
            assert x_tensor.shape[1] == truncated_bptt_steps, "tbptt split Tensor failed"

            y_tensor = torch.tensor(y_list, dtype=x_tensor.dtype)
            assert y_tensor.shape[1] == truncated_bptt_steps, "tbptt split list failed"

            pred = self(x_tensor.view(batch_size, truncated_bptt_steps))
            loss_val = torch.nn.functional.mse_loss(pred, y_tensor.view(batch_size, truncated_bptt_steps))
            return {
                "loss": loss_val,
                "hiddens": self.test_hidden,
            }

        def training_epoch_end(self, training_step_outputs):
            training_step_outputs = training_step_outputs[0]
            assert len(training_step_outputs) == (sequence_size / truncated_bptt_steps)
            loss = torch.stack([x["loss"] for x in training_step_outputs]).mean()
            self.log("train_loss", loss)

        def train_dataloader(self):
            return torch.utils.data.DataLoader(
                dataset=MockSeq2SeqDataset(),
                batch_size=batch_size,
                shuffle=False,
                sampler=None,
            )

    model = BpttTestModel(batch_size=batch_size,
                          in_features=truncated_bptt_steps, out_features=truncated_bptt_steps)
    model.example_input_array = torch.randn(5, truncated_bptt_steps)

    # fit model
    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        truncated_bptt_steps=truncated_bptt_steps,
        limit_val_batches=0,

Review comment (performance):
```suggestion
    model = BpttTestModel(
        batch_size=batch_size, in_features=truncated_bptt_steps, out_features=truncated_bptt_steps
    )
```

Code after:
            y_tensor = torch.tensor(y_list, dtype=x_tensor.dtype)
            assert y_tensor.shape[1] == truncated_bptt_steps, "tbptt split list failed"

            pred = self(x_tensor.view(batch_size, truncated_bptt_steps))
            loss_val = torch.nn.functional.mse_loss(pred, y_tensor.view(batch_size, truncated_bptt_steps))
            return {
                "loss": loss_val,
                "hiddens": self.test_hidden,
            }

        def training_epoch_end(self, training_step_outputs):
            training_step_outputs = training_step_outputs[0]
            assert len(training_step_outputs) == (sequence_size / truncated_bptt_steps)
            loss = torch.stack([x["loss"] for x in training_step_outputs]).mean()
            self.log("train_loss", loss)

        def train_dataloader(self):
            return torch.utils.data.DataLoader(
                dataset=MockSeq2SeqDataset(),
                batch_size=batch_size,
                shuffle=False,
                sampler=None,
            )

    model = BpttTestModel(batch_size=batch_size,
                          in_features=truncated_bptt_steps, out_features=truncated_bptt_steps)
    model.example_input_array = torch.randn(5, truncated_bptt_steps)

    # fit model
    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        truncated_bptt_steps=truncated_bptt_steps,
        limit_val_batches=0,
        weights_summary=None,
    )
    result = trainer.fit(model)

    assert result == 1, 'training failed to complete'