Commit cda15bb0 authored by Olegs NIKISINS's avatar Olegs NIKISINS
Browse files

Added doc-strings in the MLP class

parent cbe29954
......@@ -36,6 +36,19 @@ class TwoLayerMLP(nn.Module):
super(TwoLayerMLP, self).__init__()
"""
Init method.
Parameters
----------
in_features : int
Dimensionality of the input feature vectors.
n_hidden_relu : int
Number of ReLU units in the hidden layer of the MLP.
apply_sigmoid : bool
If set to ``True`` the sigmoid will be applied to the output of the
hidden FC layer. If ``False`` the sigmoid is not applied.
Default: ``True``.
"""
self.in_features = in_features
......@@ -52,6 +65,17 @@ class TwoLayerMLP(nn.Module):
def forward(self, x):
"""
The forward method.
Parameters
----------
x : :py:class:`torch.Tensor`
The batch to forward through the network. Size of the input batch
is [batch_size, 1, self.in_features].
Returns
-------
x : :py:class:`torch.Tensor`
Output of the MLP, class probability.
"""
# input is a batch of the size: [batch_size, 1, self.in_features],
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment