Method predict

axlearn/audio/decoder_asr.py:302–318

Computes logits. Args: input_batch: A dict containing: inputs: A Tensor of shape [batch_size, num_frames, dim]. paddings: A 0/1 Tensor of shape [batch_size, num_frames]. 1's represent paddings. Returns: Logits of shape [batch_

(self, input_batch: Nested[Tensor])

        )

    def predict(self, input_batch: Nested[Tensor]) -> Tensor:
        """Computes logits.

        Args:
            input_batch: A dict containing:
                inputs: A Tensor of shape [batch_size, num_frames, dim].
                paddings: A 0/1 Tensor of shape [batch_size, num_frames]. 1's represent paddings.

        Returns:
            Logits of shape [batch_size, num_frames, vocab_size]. Logits corresponding to padding
            frames will be 0's. Note that the returned logits are not proper log probabilities, i.e.
            we have not subtracted the log-partition function.
        """
        inputs = input_batch["inputs"]
        paddings: Tensor = input_batch["paddings"]
        logits = self.lm_head(inputs)
        return logits * safe_not(paddings)[..., None]

    def _loss_summaries(
        self,