MCPcopy Create free account
hub / github.com/pytorch/examples / EmbeddingTable

Class EmbeddingTable

distributed/rpc/rnn/rnn.py:38–55  ·  view source on GitHub ↗

r""" Encoding layers of the RNNModel

Source from the content-addressed store, hash-verified

36
37
38class EmbeddingTable(nn.Module):
39 r"""
40 Encoding layers of the RNNModel
41 """
42 def __init__(self, ntoken, ninp, dropout):
43 super(EmbeddingTable, self).__init__()
44 self.drop = nn.Dropout(dropout)
45 self.encoder = nn.Embedding(ntoken, ninp)
46 if torch.accelerator.is_available():
47 device = torch.accelerator.current_accelerator()
48 self.encoder = self.encoder.to(device)
49 nn.init.uniform_(self.encoder.weight, -0.1, 0.1)
50
51 def forward(self, input):
52 if torch.accelerator.is_available():
53 device = torch.accelerator.current_accelerator()
54 input = input.to(device)
55 return self.drop(self.encoder(input)).cpu()
56
57
58class Decoder(nn.Module):

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected