def forward(self, x): x = torch.relu(self.fc1(x)) x = torch.relu(self.fc2(x)) return torch.softmax(self.fc3(x), dim=-1)