     named: "val", in: localStorageDirectory, base: base, labelDict: labelDict)
 }
 
-func applyImageNetDataAugmentation(image: Image) -> Tensor<Float> {
+func applyImageNetDataAugmentation(image: Image, applyAugmentation: Bool) -> Tensor<Float> {
     // using the tensorflow imagenet demo from mlperf as reference:
     // https://github.com/mlcommons/training/blob/4f97c909f3aeaa3351da473d12eba461ace0be76/image_classification/tensorflow/official/resnet/imagenet_preprocessing.py#L94
     let imageData = image.tensor

          return batches.lazy.map {
            makeImageNetBatch(
              samples: $0, outputSize: outputSize, mean: mean, standardDeviation: standardDeviation,
              device: device)
          }
        }



func makeImageNetBatch<BatchSamples: Collection>(
  samples: BatchSamples, outputSize: Int, mean: Tensor<Float>?, standardDeviation: Tensor<Float>?,
  device: Device, applyAugmentation: Bool
) -> LabeledImage where BatchSamples.Element == (file: URL, label: Int32) {
  let images = samples.map(\.file).map { url -> Tensor<Float> in
    if url.absoluteString.range(of: "n02105855_2933.JPEG") != nil {
      // this is a png saved as a jpeg, we manually strip an extra alpha channel to start
      let image = Image(contentsOf: url).tensor.slice(lowerBounds: [0, 0, 0], sizes: [189, 213, 3])
      let colorOnlyImage = Image(image)
      return applyImageNetDataAugmentation(image: colorOnlyImage, applyAugmentation: applyAugmentation)
    } else {
      let image = Image(contentsOf: url)
      return applyImageNetDataAugmentation(image: image, applyAugmentation: applyAugmentation)
    }
  }

