migrate to PTL 2.0

Signed-off-by: Alexandra Antonova <antonova_sasha@list.ru>

examples/nlp/spellchecking_asr_customization/spellchecking_asr_customization_train.py

@@ -47,6 +47,10 @@
 
 @hydra_runner(config_path="conf", config_name="spellchecking_asr_customization_config")
 def main(cfg: DictConfig) -> None:
+ # PTL 2.0 has find_unused_parameters as False by default, so its required to set it to True
+ # when there are unused parameters like here
+ if cfg.trainer.strategy == 'ddp':
+ cfg.trainer.strategy = "ddp_find_unused_parameters_true"
  logging.info(f'Config Params: {OmegaConf.to_yaml(cfg)}')
 
  # Train the model

nemo/collections/nlp/models/spellchecking_asr_customization/spellchecking_model.py

@@ -272,14 +272,15 @@ def validation_step(self, batch, batch_idx):
  )
 
  val_loss = self.loss_fn(logits=logits, labels=labels, loss_mask=labels_mask)
+ self.validation_step_outputs.append({'val_loss': val_loss})
  return {'val_loss': val_loss}
 
- def validation_epoch_end(self, outputs):
+ def on_validation_epoch_end(self):
  """
  Called at the end of validation to aggregate outputs.
  :param outputs: list of individual outputs of each validation step.
  """
- avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
+ avg_loss = torch.stack([x['val_loss'] for x in self.validation_step_outputs]).mean()
 
  # Calculate metrics and classification report
  # Note that in our task recall = accuracy, and the recall column is the per class accuracy
@@ -300,12 +301,12 @@ def test_step(self, batch, batch_idx):
  """
  return self.validation_step(batch, batch_idx)
 
- def test_epoch_end(self, outputs):
+ def on_test_epoch_end(self):
  """
  Called at the end of test to aggregate outputs.
  :param outputs: list of individual outputs of each test step.
  """
- return self.validation_epoch_end(outputs)
+ return self.on_validation_epoch_end()
 
  # Functions for inference