[Caffe2] Caffe2 Model not working on android (Client). The model is working perfectly fine on Server

[pradeephike]

I am working on creating a caffe2 model for style transfer and running the model on android.

Below is the net I am using for the model generation :

` conv1_1 = brew.conv(model, data, 'conv1_1', 3, 32, 9, stride=1,pad=9//2,no_bias=1,weight_init=('GaussianFill',{'mean':0.0, 'std':1e-2}))
inst1_1 = brew.instance_norm(model,conv1_1,'inst1_1',32,epsilon=0.001)
relu1_1 = brew.relu(model, inst1_1, 'relu1_1')

conv2_2 = brew.conv(model, 'relu1_1', 'conv2_2',32, 64, 3, stride=2,pad=3//2,no_bias=1,weight_init=('GaussianFill',{'mean':0.0, 'std':1e-1}))
inst2_2 = brew.instance_norm(model,conv2_2,'inst2_2',64,epsilon=0.001)
relu2_2 = brew.relu(model, inst2_2, 'relu2_2')

conv3_3 = brew.conv(model, relu2_2, 'conv3_3', 64, 64, 3, stride=1, pad=3//2,no_bias=1,weight_init=('GaussianFill',{'mean':0.0, 'std':1e-2}))
inst3_3 = brew.instance_norm(model, conv3_3,'inst3_3',64,epsilon=0.001)
relu3_3 = brew.relu(model, inst3_3, 'relu3_3')

conv4_4 = brew.conv(model, relu3_3, 'conv4_4', 64, 64, 3, stride=1, pad=3//2,no_bias=1,weight_init=('GaussianFill',{'mean':0.0, 'std':1e-2}))
inst4_4 = brew.instance_norm(model, conv4_4,'inst4_4',64,epsilon=0.001)
residual1 = brew.sum(model,[inst4_4,relu2_2],'residual1')

conv5_5 = brew.conv(model, residual1, 'conv5_5', 64, 64, 3, stride=1, pad=3//2,no_bias=1,weight_init=('GaussianFill',{'mean':0.0, 'std':1e-2}))
inst5_5 = brew.instance_norm(model, conv5_5,'inst5_5',64,epsilon=0.001)
relu5_4 = brew.relu(model, inst5_5, 'relu5_4')

conv6_6 = brew.conv(model, relu5_4, 'conv6_6', 64, 64, 3, stride=1, pad=3//2,no_bias=1,weight_init=('GaussianFill',{'mean':0.0, 'std':1e-2}))
inst6_6 = brew.instance_norm(model, conv6_6,'inst6_6',64,epsilon=0.001)
residual2 = brew.sum(model,[inst6_6,residual1],'residual2')

deconv7_1 = brew.conv_transpose(model, residual2, 'deconv7_1', 64, 32, 3, stride=2,pad=3//2,weight_init=('GaussianFill',{'mean':0.0, 'std':1e-2}))
inst7_7 = brew.instance_norm(model, deconv7_1,'inst7_7',32,epsilon=0.001)
relu7_5 = brew.relu(model, inst7_7, 'relu7_5')

conv8_7 = brew.conv(model, relu7_5, 'conv8_7', 32, 3, 9,stride=1, pad=9//2,no_bias=1,weight_init=('GaussianFill',{'mean':0.0, 'std':1e-2}))
inst8_8 = brew.instance_norm(model,conv8_7,'inst8_8',3,epsilon=0.001)`

Now once I generate the InitNet.pb and PredictNet.pb files using mobile exporter, the model is working perfectly fine while I run inference on server side.

But once I integrate the model on the android application it gets stuck in the predictor->run.

One interesting observation - If the model was changed by removing few of the layers it is working perfectly fine on android as well.

However there are no operator difference between the two models. So logically the first model should have worked fine.

Does anyone have any idea how to tackle this problem? Because I need the first model to work on android.

P.S - If anyone needs to have a look at the model then I can share the .pb files if needed