ljeonjko:

x already has (lat, lon) dimensions.


In that case I think you can just unsqueeze dim0 and dim1 as x should already contain the values to create the batch and channel dimension.
This would create a tensor in the shape [1, 1, lat, lon] which would be an “image-like” tensor.