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