Loss key needs to be present Pytorch
I am following this repo:
https://github.com/NVIDIA/NeMo/tree/main/examples/nlp/entity_linking
Here is a small tutorial:
https://colab.research.google.com/github/NVIDIA/NeMo/blob/v1.0.2/tutorials/nlp/Entity_Linking_Medical.ipynb
Before starting this tutorial change branch to r1.10.0
When I train this model on entire UMLS dataset given the commands it gives the following error:
In automatic_optimization, when `training_step` returns a dict, the 'loss' key needs to be present

I checked the training steps method and it is fine:
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 = None
            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)

        return {&quot;loss&quot;: train_loss, &quot;lr&quot;: lr}

Here is a full stacktrace:
[NeMo I 2022-07-29 18:29:27 multi_similarity_loss:91] Encountered zero loss in multisimloss, loss = 0.0. No hard examples found in the batch
Error executing job with overrides: ['project_dir=.']
Traceback (most recent call last):
  File &quot;self_alignment_pretraining.py&quot;, line 38, in main
    trainer.fit(model)
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py&quot;, line 769, in fit
    self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/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/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py&quot;, line 93, in launch
    return function(*args, **kwargs)
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/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/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py&quot;, line 1234, in _run
    results = self._run_stage()
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py&quot;, line 1321, in _run_stage
    return self._run_train()
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py&quot;, line 1351, in _run_train
    self.fit_loop.run()
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/loops/base.py&quot;, line 204, in run
    self.advance(*args, **kwargs)
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/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/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/loops/base.py&quot;, line 204, in run
    self.advance(*args, **kwargs)
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/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/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/loops/base.py&quot;, line 204, in run
    self.advance(*args, **kwargs)
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/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/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/loops/base.py&quot;, line 204, in run
    self.advance(*args, **kwargs)
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py&quot;, line 207, in advance
    self.optimizer_idx,
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/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/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py&quot;, line 378, in _optimizer_step
    using_lbfgs=is_lbfgs,
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py&quot;, line 1593, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/core/lightning.py&quot;, line 1644, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/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/umair/miniconda3/envs/aemap/lib/python3.7/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/umair/miniconda3/envs/aemap/lib/python3.7/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/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/plugins/precision/native_amp.py&quot;, line 85, in optimizer_step
    closure_result = closure()
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py&quot;, line 148, in __call__
    self._result = self.closure(*args, **kwargs)
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py&quot;, line 134, in closure
    step_output = self._step_fn()
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py&quot;, line 437, in _training_step
    training_step_output, self.trainer.accumulate_grad_batches
  File &quot;/home/umair/miniconda3/envs/aemap/lib/python3.7/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py&quot;, line 75, in from_training_step_output
    &quot;In automatic_optimization, when `training_step` returns a dict, the 'loss' key needs to be present&quot;
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.


You get this error message about &quot;loss key needs to be present&quot; because in some training steps you return the dict {&quot;loss&quot;: None}. This happens in your code here
if train_loss == 0:
    train_loss = None
    lr = None

where you set train_loss = None. Lightning does not like that, because it wants loss to be a tensor with a graph attached.
If you wish to skip the optimization step completely, just return None from the training_step method, like this:
if train_loss == 0:
    return None

