Change final value network activation to identity

2021-08-06 18:48:53 +02:00
parent 66bfba3986
commit 025c71767f
2 changed files with 4 additions and 2 deletions
--- a/src/nets/util.py
+++ b/src/nets/util.py
@@ -1,5 +1,5 @@
 import torch
-from torch.nn import functional
+from torch.nn import functional, Identity

 def parse_functional(functional_config):
    if isinstance(functional_config, str):
@@ -9,4 +9,6 @@ def parse_functional(functional_config):
            return torch.sigmoid
        elif functional_config == 'softmax':
            return functional.softmax
+        elif functional_config == 'id':
+            return Identity()
    return None