I can use nn.Embedding trivially to get the positive embeddings using a tensor of indices like so
n = nn.Embedding(256,128)
indices = torch.randint(0, 256, (10,))
positive_embeddings = n(indices)
print(positive_embeddings.shape) #prints torch.Size([10, 128])

How do i now get all the other embeddings?
I would like to do something like
negative_embeddings = some_function(n,indices)
print(negative_embeddings.shape) #prints torch.Size([10, 255, 128])

My motivation is to use nn.CosineEmbeddingLoss with the predicted embeddings of some CNN and both the positive and negative embeddings of some target indices.