From 855dd365b539b27b818e1d42dd1b1f7102dfe5ad Mon Sep 17 00:00:00 2001 From: Eren Golge Date: Wed, 2 May 2018 04:59:23 -0700 Subject: [PATCH] bug fix --- train.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/train.py b/train.py index b529b92d..3c9e0461 100644 --- a/train.py +++ b/train.py @@ -93,6 +93,9 @@ def train(model, criterion, criterion_st, data_loader, optimizer, epoch): params_group['lr'] = current_lr optimizer.zero_grad() + + stop_target = stop_target.view(c.batch_size, stop_target.size(1) // c.r, -1) + stop_target = (stop_target.sum(2) > 0.0).long() # dispatch data to GPU if use_cuda: @@ -102,9 +105,6 @@ def train(model, criterion, criterion_st, data_loader, optimizer, epoch): linear_spec = linear_spec.cuda() stop_target = stop_target.cuda() - stop_target = stop_target.view(c.batch_size, stop_target.size(1) // c.r, -1) - stop_target = (stop_target.sum(2) > 0.0).long() - # create attention mask if c.mk > 0.0: N = text_input.shape[1]