Does anyone know if I have a tensor of dimention 20x1x256x256 how I can use make_grid and save_image to save the images in a specific color map (e.g. viridis)? 20 represents 20 different predictions and 1x256x256 represents the prediction.
I’m not sure those functions support using a custom color map (it sounds like you want to generate a heatmap first). In this case you might consider producing a heatmap (e.g., with something like matplotlib or altair) before saving the prediction images.