Upgrading to PyTorch Nightly fixed the errors and the output is looking much better now!
Edit: I spoke a bit too soon… there is also something wrong with the calibration. It seems that calibration is actually hurting performance. Here is the output of the quantized model after 1,10,100 and 1000 calibration images:
calibration1600×394 318 KB
This is what it should look like:
(output before quantization)

I have tried setting the number of CPU threads with org.pytorch.PyTorchAndroid.setNumThreads(1); but it does not make a difference. I have also tried 1,2,3,4. Is this the correct way to set the thread count?