Thanks for the reply.
Although the sizes seem to be correct. ‘embed_size’ is 256, and length of vocabulary is 256(taking values from 0 to 255) . Therefore, ‘self.embed’ should embed to 256-dimensional vector. The out of ‘linear’ is [1,256] and that of the embedding should be [1,5332,256].  In the concat, I do features.unsqueeze which will introduce one more dimension and hence the dimensions should not be a problem.
On further inspection, I realize that when I try to print embedding, right after it executes it, it gives the same error. So, somehow, embedding is getting messed up, although it does executes that statement correctly.