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
[49512]
and i want the weight matrix to be [512
128]
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 torch.mm to multiply weight by input to get output.