diff --git a/transformers_summarization_wandb.ipynb b/transformers_summarization_wandb.ipynb
index 2e16982..ff7307c 100644
--- a/transformers_summarization_wandb.ipynb
+++ b/transformers_summarization_wandb.ipynb
@@ -404,7 +404,7 @@
     "        ids = data['source_ids'].to(device, dtype = torch.long)\n",
     "        mask = data['source_mask'].to(device, dtype = torch.long)\n",
     "\n",
-    "        outputs = model(input_ids = ids, attention_mask = mask, decoder_input_ids=y_ids, lm_labels=lm_labels)\n",
+    "        outputs = model(input_ids = ids, attention_mask = mask, decoder_input_ids=y_ids, labels=labels)\n",
     "        loss = outputs[0]\n",
     "        \n",
     "        if _%10 == 0:\n",
@@ -1680,4 +1680,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 1
-}
\ No newline at end of file
+}