A softmax output

nschaetti · Jan 25, 2019 · b51fba1 · b51fba1
1 parent 2b472c8
commit b51fba1
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/echotorch/nn/ESN.py b/echotorch/nn/ESN.py
@@ -41,7 +41,8 @@ class ESN(nn.Module):
     def __init__(self, input_dim, hidden_dim, output_dim, spectral_radius=0.9, bias_scaling=0, input_scaling=1.0,
                  w=None, w_in=None, w_bias=None, w_fdb=None, sparsity=None, input_set=[1.0, -1.0], w_sparsity=None,
                  nonlin_func=torch.tanh, learning_algo='inv', ridge_param=0.0, create_cell=True,
-                 feedbacks=False, with_bias=True, wfdb_sparsity=None, normalize_feedbacks=False):
+                 feedbacks=False, with_bias=True, wfdb_sparsity=None, normalize_feedbacks=False,
+                 softmax_output=False, seed=None):
         """
         Constructor
         :param input_dim: Inputs dimension.
@@ -72,11 +73,11 @@ def __init__(self, input_dim, hidden_dim, output_dim, spectral_radius=0.9, bias_
         if create_cell:
             self.esn_cell = ESNCell(input_dim, hidden_dim, spectral_radius, bias_scaling, input_scaling, w, w_in,
                                     w_bias, w_fdb, sparsity, input_set, w_sparsity, nonlin_func, feedbacks, output_dim,
-                                    wfdb_sparsity, normalize_feedbacks)
+                                    wfdb_sparsity, normalize_feedbacks, seed)
         # end if
 
         # Ouput layer
-        self.output = RRCell(hidden_dim, output_dim, ridge_param, feedbacks, with_bias, learning_algo)
+        self.output = RRCell(hidden_dim, output_dim, ridge_param, feedbacks, with_bias, learning_algo, softmax_output)
     # end __init__
 
     ###############################################