bug fix

2018-05-02 05:00:36 -07:00 · 2018-05-02 05:00:36 -07:00 · 6b937d6059
parent 59f8d246c9
commit 6b937d6059
1 changed files with 1 additions and 1 deletions
--- 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: