@ethancaballero my first thought was that it would be nearly the same cost, because they share the same storage. So I kept numpy as the main solution for storing an image array.
However we have to keep in mind that the final object will be a torch tensor, so I compared what would be the difference of using numpy or torch, it turns out the first solution took on average 3.3 seconds to build a 512 batch and the second solution (torch) took less than 0.005 seconds. It was measured after few iterations to make sure my 6 workers were really efficient.
This speedup is the reason I decided to switch to torch tensor only. As I said, in the end you will only need a torch tensor, no value to use numpy just as a storage. I think the conversions numpy -> torch are the cause for this slowdown.
On the other hand I was pretty satisfied with the memory footprint of numpy, I hope it will stay constant with torch.