Returned tensor image without rank

Thank you for this repo, the tool is really useful!

My concern is: is it possible that tfmpl returns an image tensor with rank/shape? I want to use the returned tensor of a plot as input for a Conv layer, and using tfmpl would allow to compute the image during training phase, but the returned images have not rank

thx!