(self, x, y)
| 105 | self.cross_entropy_loss = torch.nn.CrossEntropyLoss() |
| 106 | |
| 107 | def forward(self, x, y): |
| 108 | hidden_dim = self.linear1(x) |
| 109 | output, _, _ = self.moe_1(hidden_dim) |
| 110 | output = self.linear2(output) |
| 111 | output, _, _ = self.moe_2(output) |
| 112 | output = self.linear3(output) |
| 113 | hidden_dim = hidden_dim + output |
| 114 | sentence_embed = hidden_dim.mean(1) |
| 115 | return self.cross_entropy_loss(sentence_embed, y) |
| 116 | |
| 117 | |
| 118 | class SimplePRMoEModel(torch.nn.Module): |