title: Why does Pytorch throw no loss found error even though loss key is present
body: <p>I am trying to train this model:
  <a href="https://github.com/NVIDIA/NeMo/tree/main/examples/nlp/entity_linking" rel="nofollow noreferrer">https://github.com/NVIDIA/NeMo/tree/main/examples/nlp/entity_linking</a></p>
  <p>This is their tutorial :
  <a href="https://colab.research.google.com/github/NVIDIA/NeMo/blob/v1.0.2/tutorials/nlp/Entity_Linking_Medical.ipynb" rel="nofollow noreferrer">https://colab.research.google.com/github/NVIDIA/NeMo/blob/v1.0.2/tutorials/nlp/Entity_Linking_Medical.ipynb</a></p>
  <p>Before running this make sure you are cloning from branch BRANCH = 'r1.9.0'</p>
  <p>I am training it on UMLs and following steps at the end of this notebook. The issue is that loss key is defined like this:</p>
  <pre><code>def training_step(self, batch, batch_idx):
          &quot;&quot;&quot;
          Lightning calls this inside the training loop with the data from the training dataloader
          passed in as `batch`.
          &quot;&quot;&quot;
          input_ids, token_type_ids, attention_mask, concept_ids = batch
          logits = self.forward(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
          train_loss = self.loss(logits=logits, labels=concept_ids)
  
          # No hard examples found in batch,
          # shouldn't use this batch to update model weights
          if train_loss == 0:
              train_loss = 0
              lr = None
  
          else:
              lr = self._optimizer.param_groups[0][&quot;lr&quot;]
              self.log(&quot;train_loss&quot;, train_loss)
              self.log(&quot;lr&quot;, lr, prog_bar=True)
          self.log(&quot;I do have losssssss&quot;)
          return {&quot;loss&quot;: train_loss, &quot;lr&quot;: lr}
  </code></pre>
  <p>However, it still throws the error that when training_step` returns a dict, the 'loss' key needs to be present</p>
  <p>I am using python 3.9. Running it on Ubuntu. I have installed apex like this:</p>
  <pre><code>https://github.com/NVIDIA/apex
  </code></pre>
  <p>Error occurs after first epoch.
  What could be the reason?</p>
  <p>Edit:
  Full stacktrace
  Error executing job with overrides:</p>
  <pre><code>['project_dir=./nemo/examples/nlp/entity_linking']
  Traceback (most recent call last):
    File &quot;/home/ml-002/Fatima/./nemo/examples/nlp/entity_linking/self_alignment_pretraining.py&quot;, line 38, in main
      trainer.fit(model)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py&quot;, line 768, in fit
      self._call_and_handle_interrupt(
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py&quot;, line 719, in _call_and_handle_interrupt
      return self.strategy.launcher.launch(trainer_fn, *args, trainer=self, **kwargs)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py&quot;, line 93, in launch
      return function(*args, **kwargs)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py&quot;, line 809, in _fit_impl
      results = self._run(model, ckpt_path=self.ckpt_path)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py&quot;, line 1234, in _run
      results = self._run_stage()
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py&quot;, line 1321, in _run_stage
      return self._run_train()
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py&quot;, line 1351, in _run_train
      self.fit_loop.run()
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/loops/base.py&quot;, line 204, in run
      self.advance(*args, **kwargs)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/loops/fit_loop.py&quot;, line 268, in advance
      self._outputs = self.epoch_loop.run(self._data_fetcher)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/loops/base.py&quot;, line 204, in run
      self.advance(*args, **kwargs)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py&quot;, line 208, in advance
      batch_output = self.batch_loop.run(batch, batch_idx)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/loops/base.py&quot;, line 204, in run
      self.advance(*args, **kwargs)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py&quot;, line 88, in advance
      outputs = self.optimizer_loop.run(split_batch, optimizers, batch_idx)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/loops/base.py&quot;, line 204, in run
      self.advance(*args, **kwargs)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py&quot;, line 203, in advance
      result = self._run_optimization(
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py&quot;, line 256, in _run_optimization
      self._optimizer_step(optimizer, opt_idx, batch_idx, closure)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py&quot;, line 369, in _optimizer_step
      self.trainer._call_lightning_module_hook(
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py&quot;, line 1593, in _call_lightning_module_hook
      output = fn(*args, **kwargs)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/core/lightning.py&quot;, line 1644, in optimizer_step
      optimizer.step(closure=optimizer_closure)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/core/optimizer.py&quot;, line 168, in step
      step_output = self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/strategies/ddp.py&quot;, line 278, in optimizer_step
      optimizer_output = super().optimizer_step(optimizer, opt_idx, closure, model, **kwargs)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/strategies/strategy.py&quot;, line 193, in optimizer_step
      return self.precision_plugin.optimizer_step(model, optimizer, opt_idx, closure, **kwargs)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/plugins/precision/native_amp.py&quot;, line 85, in optimizer_step
      closure_result = closure()
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py&quot;, line 148, in __call__
      self._result = self.closure(*args, **kwargs)
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py&quot;, line 134, in closure
      step_output = self._step_fn()
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py&quot;, line 436, in _training_step
      result = self.output_result_cls.from_training_step_output(
    File &quot;/home/ml-002/anaconda3/envs/py39_ft/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py&quot;, line 74, in from_training_step_output
      raise MisconfigurationException(
  pytorch_lightning.utilities.exceptions.MisconfigurationException: In automatic_optimization, when `training_step` returns a dict, the 'loss' key needs to be present
  
  Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
  Epoch 0:   1%| | 133533/16962478 [2:28:05&lt;311:03:01, 15.03it/s, loss=0.106, v_nu
  </code></pre>