Missing persist on trainDataset in ml.classification.GBTClassifier.train()

  override protected def train(
      dataset: Dataset[_]): GBTClassificationModel = instrumented { instr =>
    val categoricalFeatures: Map[Int, Int] =
      MetadataUtils.getCategoricalFeatures(dataset.schema($(featuresCol)))
    ...
    val numFeatures = trainDataset.first().features.size // first use trainDataset
    ...
    // trainDataset will be used by other actions in run methods.    
    val (baseLearners, learnerWeights) = if (withValidation) {
      GradientBoostedTrees.runWithValidation(trainDataset, validationDataset, boostingStrategy,
        $(seed), $(featureSubsetStrategy))
    } else {
      GradientBoostedTrees.run(trainDataset, boostingStrategy, $(seed), $(featureSubsetStrategy))
    }