Comment StopNet arguments

This commit is contained in:
Eren Golge 2018-05-23 06:18:27 -07:00
parent 14f9d06b31
commit 8ffc85008a
1 changed files with 7 additions and 1 deletions

View File

@ -331,7 +331,13 @@ class Decoder(nn.Module):
class StopNet(nn.Module):
def __init__(self, r, memory_dim):
"""Predicts the stop token to stop the decoder at testing time"""
r"""
Predicts the stop token to stop the decoder at testing time
Args:
r (int): number of network output frames.
memory_dim (int): single feature dim of a single network output frame.
"""
super(StopNet, self).__init__()
self.rnn = nn.GRUCell(memory_dim * r, memory_dim * r)
self.relu = nn.ReLU()