From 43b19ab1b359e4768f4ff23c48bd6c436bb69157 Mon Sep 17 00:00:00 2001 From: Eren Golge Date: Sat, 5 May 2018 06:34:49 -0700 Subject: [PATCH] bug fix --- train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/train.py b/train.py index c33b77df..34a58d57 100644 --- a/train.py +++ b/train.py @@ -95,7 +95,7 @@ def train(model, criterion, criterion_st, data_loader, optimizer, epoch): optimizer.zero_grad() stop_target = stop_target.view(text_input.shape[0], stop_target.size(1) // c.r, -1) - stop_target = (stop_target.sum(2) > 0.0).float().unsqueeze(2) + stop_target = (stop_target.sum(2) > 0.0).float() # dispatch data to GPU if use_cuda: @@ -240,7 +240,7 @@ def evaluate(model, criterion, criterion_st, data_loader, current_step): stop_target = data[5] stop_target = stop_target.view(text_input.shape[0], stop_target.size(1) // c.r, -1) - stop_target = (stop_target.sum(2) > 0.0).float().unsqueeze(2) + stop_target = (stop_target.sum(2) > 0.0).float() # dispatch data to GPU if use_cuda: