How to make Colormaps for 2d feature maps

1

I am trying to reconstruct a neural net for shape and appearance disentanglement. To track progress, I want to produce some color maps that look as presented in the paper:

enter image description here

The maps have the following shape: [batch_size, n_maps, 64, 64].

n_maps equals 16, so I basically want to plot 16 different color maps. Is there an elegant way to do that?