The input tensor should have dimensions 1 x height x width x 3. Got 1 x 3 x 224 x 224
I want to convert the Pytorch-trained model to the tensorflow model and use the model on mobile devices. For this, I follow these steps; First I convert the pytorch trained model to onnx format. Then I convert the onnx format to the tensorflow model.
Firstly pytorch trained model to onnx;
import torch
import torch.onnx
from detectron2.modeling import build_model
from detectron2.modeling import build_backbone
from torch.autograd import Variable

model= build_backbone(cfg)
model.eval()

dummy_input = torch.randn(1,3,224, 224,requires_grad=True)

torch.onnx.export(model,dummy_input,&quot;drive/Detectron2/model_final.onnx&quot;)

Then onnx to tflite model;
import onnx
import warnings
from onnx_tf.backend import prepare

model = onnx.load(&quot;drive/Detectron2/model_final.onnx&quot;)
tf_rep = prepare(model)
tf_rep.export_graph(&quot;drive/Detectron2/tf_rep.pb&quot;)

import tensorflow as tf

## TFLite Conversion
# Before conversion, fix the model input size
model = tf.saved_model.load(&quot;drive/Detectron2/tf_rep.pb&quot;)
model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY].inputs[0].set_shape([1, 3,224, 224])
tf.saved_model.save(model, &quot;saved_model_updated&quot;, signatures=model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY])
# Convert
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir='saved_model_updated', signature_keys=['serving_default'])
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]
tflite_model = converter.convert()

# Save the model.
with open('drive/Detectron2/model.tflite', 'wb') as f:
  f.write(tflite_model)

## TFLite Interpreter to check input shape
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# Test the model on random input data.
input_shape = input_details[0]['shape']
print(input_shape)

But when I use the model on mobile devices, I get the following error;
java.lang.AssertionError: Error occurred when initializing ImageSegmenter: The input tensor should have dimensions 1 x height x width x 3. Got 1 x 3 x 224 x 224.

Where am I doing wrong?

Maybe you could try einops for tensor transformations. It's elegant and powerful.
In your case, the code should be

import einops
input_tensor = einops.rearrange(input_tensor,'b c w h -&gt; b w h c')

