@morpheusPrime I am not sure if pytorch have such a function. I didn’t check. However I was thinking the output of your model will be a batch of say 4 image then I guess you have to define a method for yourself in which you do the required post-processing as I did in my code and you need to use the Process Pool to call your post-processing method each time you predict. The reason is that during preprocessing we had the entire dataset at one place so we called the Process Pool only once and it helped but now we will be getting the output in batches.