def METHOD_NAME(self, batch, batch_idx):
    """
    Lightning calls this inside the training loop with the data from the training dataloader
    passed in as `batch`.
    """
    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]["lr"]
        self.log("train_loss", train_loss)
        self.log("lr", lr, prog_bar=True)
    return {"loss": train_loss, "lr": lr}