Custom Yolov5s TorchScript model does not work on Android Object Detection Demo App

I have custom model trained on yolov5s v5 and I converted it to torchscript.ptl using ultralytics with code modification as told here.
When I build on android i get the following error:

error that appears in debugging:
E/AndroidRuntime: FATAL EXCEPTION: Thread-2
Process: org.pytorch.demo.objectdetection, PID: 27891
java.lang.ArrayIndexOutOfBoundsException: length=1310400; index=1310449
at org.pytorch.demo.objectdetection.PrePostProcessor.outputsToNMSPredictions(

I changed my torch version also, I am still facing the same issue.

yolov5s.torchscript.ptl model file provided in the repo. works fine. I’m facing this issue when I use my custom model.

cc: @IvanKobzarev @Linbin

there are some predefined parameters: android-demo-app/ at 76ba0e04a423acbfe960dcd8dd9a0bc47c3893e7 · pytorch/android-demo-app · GitHub

Please check if it fits your new model.

Thank you for your reply. I changed the value as calculated by 25200*(num_of_class+5) but I’m still facing the same issue.
Do I need to change this value also or any other ?

“private static int mOutputColumn = 85; // left, top, right, bottom, score and 80 class probability”

I changed only " private static int mOutputRow = 25200; // as decided by the YOLOv5 model for input image of size 640*640" this value.

I even changed mOutputColumn = 85 to my custom class number and class number + 5 also but I am still facing the same issue.