Following this, a custom initialization results in a test accuracy of 76%:
nn.init.uniform_(self.embedding.weight, -0.05, 0.05)
nn.init.xavier_uniform_(self.fc.weight, 1.0)
nn.init.zeros_(self.fc.bias)
1 Like