predict method

Map<String, Tensor> predict(
  1. Tensor state
)

Implementation

Map<String, Tensor> predict(Tensor state) {
  Tensor logits = transformer.lmHead.forward(state);
  Tensor vHidden = state.matmul(transformer.vW1);
  for (int i = 0; i < vHidden.data.length; i++) if (vHidden.data[i] < 0) vHidden.data[i] = 0;
  Tensor value = vHidden.matmul(transformer.vW2);
  value.data[0] = _mathTanh(value.data[0]);

  return {"policy": logits, "value": value};
}