Unable to call metric from any step in Lightning module about torchmetrics HOT 9 CLOSED

class MyModel(LightningModule):
    def __init__(self, config: DictConfig, vocabulary: Vocabulary):
        super().__init__()
        ...
        pad_idx = vocabulary.label_to_id[PAD]
        ignore_idx = [vocabulary.label_to_id[i] for i in [UNK, EOS, SOS] if i in vocabulary.label_to_id]
        self._train_metrics = SequentialF1Score(mask_after_pad=True, pad_idx=pad_idx, ignore_idx=ignore_idx)
        self._val_metrics = SequentialF1Score(
            mask_after_pad=True, pad_idx=pad_idx, ignore_idx=ignore_idx, compute_on_step=False
        )
        self._test_metrics = SequentialF1Score(
            mask_after_pad=True, pad_idx=pad_idx, ignore_idx=ignore_idx, compute_on_step=False
        )
        ....

    ...

    def training_step(self, batch: Tuple[torch.Tensor, dgl.DGLGraph], batch_idx: int) -> Dict:  # type: ignore
        labels, graph = batch
        # [seq length; batch size; vocab size]
        logits = self(graph, labels.shape[0], labels)
        loss = self._calculate_loss(logits, labels)
        prediction = logits.argmax(-1)

        batch_metrics: ClassificationMetrics = self._train_metrics(prediction, labels)  # **Error happens here**
        log = {
            "train/loss": loss,
            "train/f1": batch_metrics.f1_score,
            "train/precision": batch_metrics.precision,
            "train/recall": batch_metrics.recall,
        }
        self.log_dict(log)
        self.log("f1", batch_metrics.f1_score, prog_bar=True, logger=False)

        return {"loss": loss}

    def validation_step(self, batch: Tuple[torch.Tensor, dgl.DGLGraph], batch_idx: int) -> Dict:  # type: ignore
        labels, graph = batch
        # [seq length; batch size; vocab size]
        logits = self(graph, labels.shape[0], labels)
        loss = self._calculate_loss(logits, labels)
        prediction = logits.argmax(-1)

        self._val_metrics(prediction, labels)  # **And here**
        return {"loss": loss}

wandb_logger.watch(model)