You can input a batch when testing like this:
at::Tensor input1 = torch::ones({1,3,128,128});
at::Tensor input2 = torch::ones({1,3,128,128});
at::Tensor inputs = torch::cat({input1,input2},0);
at::Tensor output = module->forward({inputs}).toTensor();