Running predictions using CoreML sometimes fails

I've converted and am trying to run a model that does semantic segmentation on an image. The issue I'm running into is that the model sometimes returns results as expected, and sometimes returns an MLMultiArray in the right shape, filled with NaNs. Even if I run on the same input, without changing a single line of code, I'll sometimes get the correct results and other times NaNs. Even if there was some bug in the model/conversion, I'd still expect it to produce results, and that those results would be relatively consistent between runs. Am I missing something here? Is there any way to get more insight into what's going on inside and what's the problem? When the model does return results, the results seem to make sense.

I published my code here: https://github.com/Ofir-Purple/ios-coreml-segmentation-tests , this repo includes the XCode project, as well as explanations on how to download the mlmodel file, or convert it anew from Caffe (it's based on VGG16, so it's 500+MB, I'm planning to replace it with another smaller version, which I'm currently working on).

The converted model is based on the Caffe model and code found here: https://github.com/fyu/dilation

The model contains dilated convolutions, which I saw might cause problems on older iOS devices, but AFAICT should be okay when run in the simulator.

The code that's doing the predictions:

func predictUsingCoreML(image: UIImage) {
    
    let model = dilation8()
    
    // The image is already after the pre-processing steps (done in Python).
    // TODO: Implement the pre-processing steps here.
    
    // Get the image's CVPixelBuffer
    let pixelBuffer = image.pixelBuffer(width: 900, height: 900)!
    
    guard let output = try? model.prediction(data: pixelBuffer) else {
        print("failed")
        return
    }
    
    // output.prob should be a 21*66*66 array (classes_probs*height*width)
    // Sometimes I get back a 21*66*66 array filled with NaN's and sometimes with probs :(
    var currProb: Double?
    for index in 0...output.prob.count - 1 {
        
        currProb = output.prob[index] as? Double
        if !(currProb?.isNaN)! {
            print("\(index) : \(String(describing: currProb))")
        }
    }
    
    // TODO: Post-processing...
    
    print("Done!")
}

The loop at the end sometimes prints the predictions, and sometimes nothing (because output.prob is filled with NaNs).

Thanks!