It seems you want to implement the CBOW setup of Word2Vec. You can easily find PyTorch implementations for that. For example, I found this implementation 497 in 10 seconds :).
This example uses nn.Embedding so the inputs of the forward() method is a list of word indexes (the implementation doesn’t seem to use batches). But yes, instead of nn.Embedding you could use nn.Linear. The only change needed would be that inputs not has to be a list of one-hot vectors. But I wouldn’t bother, nn.Embedding keeps things simpler.