Method predict

axlearn/common/trainer_test.py:254–258

(self, input_batch: NestedTensor)

        label: Tensor = input_batch["label"]
        loss = (
            -(jax.nn.log_softmax(logits) * jax.nn.one_hot(label, NUM_CLASSES, dtype=logits.dtype))
            .sum(axis=-1)
            .mean()
        )
        return loss, {"prng_key": self.prng_key}

    def predict(self, input_batch: NestedTensor) -> Tensor:
        image = input_batch["image"]