# How to use make\_grid to display the tensors in color map?

**URL:** <https://discuss.pytorch.org/t/how-to-use-make-grid-to-display-the-tensors-in-color-map/123508>\
**Category:** vision\
**Created:** [June 7, 2021, 2:38pm UTC](https://discuss.pytorch.org/t/how-to-use-make-grid-to-display-the-tensors-in-color-map/123508 "2021-06-07T14:38:36Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![seyeeet](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/seyeeet/32/21931_2.png) [@seyeeet](https://discuss.pytorch.org/u/seyeeet)\
**Post date:** [June 7, 2021, 2:38pm UTC](https://discuss.pytorch.org/t/how-to-use-make-grid-to-display-the-tensors-in-color-map/123508/1 "2021-06-07T14:38:36Z")

</div>

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.

---

<div class="post-metadata">

**Author:** ![eqy](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/eqy/32/46233_2.png) [@eqy](https://discuss.pytorch.org/u/eqy)\
**Post date:** [June 7, 2021, 7:10pm UTC](https://discuss.pytorch.org/t/how-to-use-make-grid-to-display-the-tensors-in-color-map/123508/2 "2021-06-07T19:10:03Z")

</div>

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.