Core ML on iOS 15 consumes much larger memory for prediction

Hi,

I found that the Core ML on iOS 15 consumes almost 450MB larger memory than iOS 14 for prediction of my model. It may easily cause running out of memory crash on small memory devices. To reproduce the issue, I created a sample code as bellow. It shows around 450MB on iOS 15 but only about 4MB on iOS 14.X. I have already submitted the issue through Feedback Assistance.
But does anyone have any idea to work around the issue? I've confirmed the the issue still exists on iOS 15.1 Beta 2.

Any advice would be appreciated.

import CoreML

class ViewController: UIViewController {

  @IBOutlet weak var memoryLabel: UILabel!
  @IBAction func predAction(_ sender: Any) {
    let modelConf = MLModelConfiguration()
    modelConf.computeUnits = .cpuOnly
    do {
      let input = UnsafeMutablePointer<Float>.allocate(capacity: 512 * 1024 * 2)
      input.initialize(repeating: 0.0, count: 512 * 1024 * 2)
      let inputMultiArray = try MLMultiArray(dataPointer: input, shape: [1,512,1024,2], dataType: .float32, strides: [(512 * 1024 * 2) as NSNumber ,(1024 * 2) as NSNumber,2,1], deallocator: nil)
      let model = try TestModel(configuration: modelConf)
      let memBefore = os_proc_available_memory()
      let _ = try model.prediction(input:TestModelInput(strided_slice_3_0: inputMultiArray))
      let memAfter = os_proc_available_memory()
      print("memory size for prediction",memBefore - memAfter)
      memoryLabel.text = String(describing:(memBefore - memAfter) )
       
    }
    catch {
      print("error")
    }
  }
  override func viewDidLoad() {
    super.viewDidLoad()
    // Do any additional setup after loading the view.
  }


}