From 6b937d6059b8dc36298ed17f66580caf8d830464 Mon Sep 17 00:00:00 2001
From: Eren Golge <egolge@mozilla.com>
Date: Wed, 2 May 2018 05:00:36 -0700
Subject: [PATCH] bug fix

---
 train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train.py b/train.py
index f7d0a273..79e12b6a 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(c.batch_size, stop_target.size(1) // c.r, -1)
-        stop_target = (stop_target.sum(2) > 0.0).float()
+        stop_target = (stop_target.sum(2) > 0.0).float().unsqueeze(2)
 
         # dispatch data to GPU
         if use_cuda: