StateIn Input for ActivityClassifier in CoreML

Question

I created ActivityClassifier using CreateML and implemented the below code to make predictions.

The resulting model has “StateIn” input parameter. And I read to use “StateOut” output parameter to feed back into StateIn for the next prediction. So I do stateIn = prediction.stateOut. When I do this, the prediction always comes out to null probability.

When I comment out stateIn = prediction.stateOut, effectively always passing empty stateIn, the prediction does come out with a valid probability. What am I doing wrong here?

// Define the PredictionOutput tuple type
typealias PredictionOutput = (label: String, confidence: Double)

class Predictor {
    let activityClassifier = try! MyActivityClassifier_4(configuration: .init())
    lazy var model = activityClassifier.model
    
    let predictionWindowSize = 30
    
    // Get the previous state output
    var stateIn: MLMultiArray?
    
    func isReadyToMakePrediction() -> Bool {
        motionDataWindow.count == predictionWindowSize
    }
    
    func makePrediction() throws -> PredictionOutput {
        let accelX = try MLMultiArray(motionDataWindow.map { $0.accel_x })
        let accelY = try MLMultiArray(motionDataWindow.map { $0.accel_y })
        let accelZ = try MLMultiArray(motionDataWindow.map { $0.accel_z })
        let gyroX = try MLMultiArray(motionDataWindow.map { $0.gyro_x })
        let gyroY = try MLMultiArray(motionDataWindow.map { $0.gyro_y })
        let gyroZ = try MLMultiArray(motionDataWindow.map { $0.gyro_z })

        if stateIn == nil {
            if let constraint = model.modelDescription.inputDescriptionsByName["stateIn"]?.multiArrayConstraint {
                stateIn = try MLMultiArray(shape: constraint.shape, dataType: constraint.dataType)
            }
        }
        
        let predictionInput = MyActivityClassifier_4Input(
            accelerometerAccelerationX_G_: accelX,
            accelerometerAccelerationY_G_: accelY,
            accelerometerAccelerationZ_G_: accelZ,
            motionRotationRateX_rad_s_: gyroX,
            motionRotationRateY_rad_s_: gyroY,
            motionRotationRateZ_rad_s_: gyroZ,
            stateIn: stateIn!
        )
        
        let prediction = try activityClassifier.prediction(input: predictionInput)
        
        let confidence = prediction.labelProbability[prediction.label]!
        
        
        // Update the state
        stateIn = prediction.stateOut.  // <-----
        
        return (
            label: prediction.label,
            confidence: prediction.labelProbability[prediction.label]!
        )
    }
}