status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
        immediate_execution_ctx_.get(), X_vals, X_dims, 2, &X_raw);
    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
    X.reset(X_raw);
  }
  // Label
  int32_t Y_vals[] = {1, 0, 1};
  int64_t Y_dims[] = {3};
  AbstractTensorHandlePtr Y;
  {
    AbstractTensorHandle* Y_raw;
    status_ = TestTensorHandleWithDims<int32_t, TF_INT32>(