{"slug": "building-lstms-with-pytorch-and-lightning-ai-part-4-training-step-and-initial", "title": "Building LSTMs with PyTorch and Lightning AI Part 4: Training Step and Initial Predictions", "summary": "A developer building LSTMs with PyTorch and Lightning AI implemented the training_step function and ran initial predictions without training. The model predicted Company A's stock price reasonably close to the observed value but Company B's prediction was far off, indicating the need for training.", "body_md": "In the [previous article](https://dev.to/rijultp/building-lstms-with-pytorch-and-lightning-ai-part-3-finishing-the-lstm-cell-1fo0), we finished the LSTM cell, explored the forward method and the Adam optimizer for the model.\n\nIn this article, we will explore the `training_step()`\n\nfunction, and try to run the model without training.\n\nThe `training_step()`\n\nfunction takes a batch of training data from one of the two companies, along with the index of that batch.\n\nIt then uses the `forward()`\n\nfunction to make a prediction for that training example.\n\n``` python\ndef training_step(self, batch, batch_idx):\n    input_i, label_i = batch\n    output_i = self.forward(input_i[0])\n    loss = (output_i - label_i)**2\n```\n\nNext, it calculates the loss, which is the squared residual between the predicted value and the observed value.\n\nWe can also log the loss to easily track how it changes during training.\n\nLightning provides the `log()`\n\nfunction for this purpose. It automatically stores the logs in a `lightning_logs`\n\ndirectory.\n\nWe can log other values as well, such as the predictions for Company A and Company B.\n\nFinally, we return the loss.\n\n``` python\ndef training_step(self, batch, batch_idx):\n    input_i, label_i = batch\n    output_i = self.forward(input_i[0])\n    loss = (output_i - label_i)**2\n\n    self.log(\"train_loss\", loss)\n\n    if label_i == 0:\n        self.log(\"out_0\", output_i)\n    else:\n        self.log(\"out_1\", output_i)\n\n    return loss\n```\n\nSo far, we have implemented the following:\n\n`lstm_unit()`\n\n.`forward()`\n\nmethod to perform a forward pass through the unrolled LSTM.`configure_optimizers()`\n\n.`training_step()`\n\n.Now let's try using the model.\n\n```\nmodel = LSTMByHand()\n\nprint(\"\\nComparing observed and predicted values\")\n\nprint(\n    \"Company A: Observed = 0, Predicted =\",\n    model(torch.tensor([0., 0.5, 0.25, 1.])).detach()\n)\n\nprint(\n    \"Company B: Observed = 1, Predicted =\",\n    model(torch.tensor([1., 0.5, 0.25, 1.])).detach()\n)\n```\n\nHere, we pass a tensor containing the stock prices for Days 1 through 4. The model then predicts the value for Day 5.\n\nThe model returns both the prediction and its associated computation graph. We call `.detach()`\n\nto remove the computation graph and retrieve only the prediction.\n\nRunning the code produces the following output:\n\n```\nComparing observed and predicted values\nCompany A: Observed = 0, Predicted = tensor(-0.2321)\nCompany B: Observed = 1, Predicted = tensor(-0.2360)\n```\n\nThe prediction for Company A is reasonably close to the observed value.\n\nHowever, the prediction for Company B is quite far from the expected value.\n\nIn the next article, we will train the model to improve these predictions.\n\nAI agents write code fast. They also silently remove logic, change behavior, and introduce bugs -- without telling you. You often find out in production.\n\n[git-lrc](https://github.com/HexmosTech/git-lrc) fixes this. It hooks into git commit and reviews every diff before it lands. 60-second setup. Completely free.\n\nAny feedback or contributors are welcome! It's online, source-available, and ready for anyone to use.\n\nGive it a ⭐ [star on Github](https://github.com/HexmosTech/git-lrc)", "url": "https://wpnews.pro/news/building-lstms-with-pytorch-and-lightning-ai-part-4-training-step-and-initial", "canonical_source": "https://dev.to/rijultp/building-lstms-with-pytorch-and-lightning-ai-part-4-training-step-and-initial-predictions-33fj", "published_at": "2026-06-26 19:20:44+00:00", "updated_at": "2026-06-26 19:34:09.766083+00:00", "lang": "en", "topics": ["machine-learning", "neural-networks", "developer-tools"], "entities": ["PyTorch", "Lightning AI", "Adam optimizer", "LSTM", "Company A", "Company B", "git-lrc", "HexmosTech"], "alternates": {"html": "https://wpnews.pro/news/building-lstms-with-pytorch-and-lightning-ai-part-4-training-step-and-initial", "markdown": "https://wpnews.pro/news/building-lstms-with-pytorch-and-lightning-ai-part-4-training-step-and-initial.md", "text": "https://wpnews.pro/news/building-lstms-with-pytorch-and-lightning-ai-part-4-training-step-and-initial.txt", "jsonld": "https://wpnews.pro/news/building-lstms-with-pytorch-and-lightning-ai-part-4-training-step-and-initial.jsonld"}}