Function InferLarsUpdateTensorDesc

oneflow/user/ops/model_update_ops.cpp:480–494

  return Maybe<void>::Ok();
}
Maybe<void> InferLarsUpdateTensorDesc(user_op::InferContext* ctx) {
  const user_op::TensorDesc& model = ctx->InputTensorDesc("model", 0);

  const Shape& shape = model.shape();
  const user_op::TensorDesc& model_diff = ctx->InputTensorDesc("model_diff", 0);
  CHECK_EQ_OR_RETURN(model_diff.shape(), shape);
  const user_op::TensorDesc& momentum = ctx->InputTensorDesc("momentum", 0);
  JUST(CheckShapeLike(&momentum, &model));
  JUST(CheckLearningRateShape(ctx));
  if (ctx->has_input("scale_by_tensor", 0)) {
    const auto& scale_by_tensor = ctx->InputTensorDesc("scale_by_tensor", 0);
    JUST(CheckScalarShape(&scale_by_tensor));
  }
  return Maybe<void>::Ok();
}
Maybe<void> InferLarsUpdateDataType(user_op::InferContext* ctx) {
  const user_op::TensorDesc& model = ctx->InputTensorDesc("model", 0);
  const user_op::TensorDesc& momentum = ctx->InputTensorDesc("momentum", 0);