Method forward

axlearn/common/inference_test.py:165–167

(self, input_batch: NestedTensor)

        self.predict_dtypes = []

    def forward(self, input_batch: NestedTensor) -> tuple[Tensor, NestedTensor]:
        y = self.predict(input_batch)
        return y.mean(), {}

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