How can i embed an image i.e build a 2-D weight matrix?

i know how to use Embedding() but it is used with sequences, but in my case the input is 2-D matrix
and i want the weight matrix to be [512
so that the result be [49*128]

create an nn.Parameter weight of your desired shape in the constructor of your model, and then in forward, just to multiply weight by input to get output.