diff options
Diffstat (limited to 'src')
| -rwxr-xr-x | src/nn_train.py | 5 |
1 files changed, 3 insertions, 2 deletions
diff --git a/src/nn_train.py b/src/nn_train.py index 8f1a73e..9e2b482 100755 --- a/src/nn_train.py +++ b/src/nn_train.py @@ -49,6 +49,7 @@ import torch from tqdm import tqdm from torch.nn import Embedding, MSELoss, Linear, Module +from torch.nn.functional import silu, tanh from torch import Tensor @@ -211,8 +212,8 @@ class Net(Module): # inputs, the rest are raw cell states (0..6) n = x.shape[0] inp = torch.cat([x[:, :3], self.embedding(x[:, 3:].long()).view(n, -1)], dim=1) - hidden = torch.relu(self.hidden(inp)) - return torch.tanh(self.output(hidden)).squeeze(-1) + hidden = silu(self.hidden(inp)) + return tanh(self.output(hidden)).squeeze(-1) def fit( self, |