How to get required input shape of traced model on libtorch

C++
1

Hi, The input size of the model needs to be estimated in libtorch.
Here’s what I tried:

Step

  1. load PyTorch model on Python
  2. save traced model on Python
  3. load traced model on libtorch (C++)

1-2, load and save on Python

import torch
import torchvision
from torchvision.models import ResNet18_Weights

model = torchvision.models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1).cuda()
model.eval()
example_in = torch.rand(1, 3, 224, 224).cuda()
traced_model = torch.jit.trace(model, example_in)
torch.jit.save(traced_model, "traced_resnet_model.pt")

3, load traced model on libtorch

#include <torch/torch.h>
#include <torch/script.h> // One-stop header.

int main(void)
{
	torch::jit::script::Module module;
	try {
		module = torch::jit::load("path/to/traced_resnet_model.pt");
	}
	catch (const c10::Error& e) {
		std::cerr << "Failed\n";
		return -1;
	}

	auto input_type = module.get_method("forward")
		.function()
		.getSchema()
		.arguments()[1]
		.type()
		->cast<c10::TensorType>();

	if (input_type) {
		auto sizes = input_type->sizes();
		if (!sizes.concrete_sizes()) {
			std::cerr << "Unable to estimate input size information." << std::endl;
		}
		else {
			std::cout << "Input tensor shape: ";
			for (const auto& size : *sizes.concrete_sizes()) {
				std::cout << size << " ";
			}
			std::cout << std::endl;
		}
	}
	else {
		std::cerr << "Unable to get TensorType information." << std::endl;
	}

	return 0;

I want to get the same dimension as example_in.
Can anyone help me.
Thank you !