Masking the Latent Dimension to make the same size as an image

How would would take latent code with dimensions (batch_size,K) and then do some element wise multiplication to turn that matrix into dim (batch_size,channels, height, width) where height and width will be N by N matrix essentially and you fill K dimensions into the height, width dimensions. Since K will be less than N by N matrix or (h,w) there will be n-k zero entries.

An elementwise multiplication won’t change the number of elements in the tensors, so you would have to view or reshape the tensor into the desired shape.
If channels*height*width is larger than K, you could create a new tensor via output = torch.zeros(batch_size, channels, height, width) and copy the reshaped input into it.