Add 'same' as a dec option for s2s, fix runtime error on GPUs
@@ -58,11 +58,11 @@ class Seq2seqAgent(Agent):
agent.add_argument('-enc', '--encoder', default='gru',
choices=['rnn', 'gru', 'lstm'],
help='Choose between different encoder modules.')
- agent.add_argument('-dec', '--decoder', default='shared',
- choices=['shared', 'rnn', 'gru', 'lstm'],
+ agent.add_argument('-dec', '--decoder', default='same',
+ choices=['same', 'shared', 'rnn', 'gru', 'lstm'],
help='Choose between different decoder modules. '
- 'If set to shared, uses the exact same module'
- ' and weights as the encoder.')
+ 'Default "same" uses same class as encoder, '
+ 'while "shared" also uses the same weights.')
def __init__(self, opt, shared=None):
"""Set up model if shared params not set, otherwise no work to do."""
@@ -126,6 +126,8 @@ class Seq2seqAgent(Agent):
# decoder produces our output states
if opt['decoder'] == 'shared':
self.decoder = self.encoder
+ elif opt['decoder'] == 'same':
+ self.decoder = enc_class(hsz, hsz, opt['numlayers'])
else:
dec_class = opt_to_class[opt['decoder']]
self.decoder = dec_class(hsz, hsz, opt['numlayers'])
@@ -205,7 +207,6 @@ class Seq2seqAgent(Agent):
self.decoder.cuda()
self.h2o.cuda()
self.dropout.cuda()
- self.softmax.cuda()
if self.use_attention:
self.attn.cuda()
self.attn_combine.cuda()
