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

      cropped = Tensor<Float>([offsetY, targetX, targetY, offsetX])
    }

    let imageBroadcast = imageData.reshaped(to: [1, height, width, channels])
    let bboxBroadcast = cropped.reshaped(to: [1, 4])



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)
    }
  }

