Classification model keeps predicting same class in android

I have a trained a classification model to classify black and white segmented outputs.
The pixels consists of only two values, 0 for black and 1 for white.
The image looks like this:

Before I run the classification model, the original color image has gone though a segmentation model to get the the black and white bitmap. The segmentation model works fine and the results are pretty similar to the PC version, but when I put the black and white bitmap into the mobile classification model, it keeps predicting the same class.

I tried exporting the black and white bitmap to PNG and ran it on the PC side, it is able to predict the correct class.

I created the black and white bitmap like this:

for (int j = 0; j < scores.length; j++) {
            if (scores[j]>0.4){
                intValues[j] = 0xFFFFFFFF; // white pixels

                intValues[j] =0xFF000000; // black pixels


I’ve optimized the model for mobile like this:

script_module = torch.jit.script(pc_classification_model)
optimized_scripted_module = optimize_for_mobile(script_module )

Library used for Pytorch Android :

implementation 'org.pytorch:pytorch_android_lite:1.10.0'
implementation 'org.pytorch:pytorch_android_torchvision_lite:1.10.0'

I think it’s related to input normalization. Different image format can lead to different input values. You may print the input tensor to debug.

This is the input for the segmentation, the segmentation does not have any problems though, only the classification. This is my classification code:

//outputBitmap is the segmented black and white image
final Tensor segmentedTensor = TensorImageUtils.bitmapToFloat32Tensor(outputBitmap,mean, std); 
final Tensor outputTensor1 = module_classification.forward(IValue.from(segmentedTensor)).toTensor();
float[] scores1 = outputTensor1.getDataAsFloatArray();

float maxScore = -Float.MAX_VALUE;
int maxScoreIdx = -1;
for (int i = 0; i < scores1.length; i++) {
    if (scores1[i] > maxScore) {
        maxScore = scores1[i];
        maxScoreIdx = i;
String className = Constants.CLASSES[maxScoreIdx];

No normalization is used, so i set the mean and std to

float[] mean = {0.0f, 0.0f, 0.0f};
float[] std = {1.0f, 1.0f, 1.0f};

I printed the input tensor for the classification too, it does not seem to be normalization problem as the input consists of only 0s and 1s, same as with PC

The same model should produce the same output given the same input on both mobile and server. It’s good that you validate that the segmentedTensor has the proper range. Does it also have the proper shape? Do you get different results from the classification model when classifying different inputs?