Libtorch Input Batch

The above is the data preprocessing stage. The batch input to the model is one.
Now, I want to set batch to four, because will process four channels of video at the same time.
How should I do ? Hope get your help. Thank you !

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();