From ccd91a67d0d5ed13906dfa2f00014ca352ec9bbf Mon Sep 17 00:00:00 2001 From: Eren Golge Date: Wed, 2 May 2018 04:56:52 -0700 Subject: [PATCH] bug fix --- train.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/train.py b/train.py index a3a93962..319c8a63 100644 --- a/train.py +++ b/train.py @@ -102,7 +102,8 @@ 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(B, stop_target.size(1) // c.r, -1) + stop_target = stop_target.view(c.batch_size, stop_target.size(1) // c.r, -1) + import ipdb; ipdb.set_trace() stop_target = (stop_target.sum(1) > 0.0).long() # create attention mask