guided attn #6

2018-04-24 11:32:22 -07:00 · 2018-04-24 11:32:22 -07:00 · 60a0721bf3
parent aea26bf55a
commit 60a0721bf3
1 changed files with 1 additions and 1 deletions
--- a/train.py
+++ b/train.py
@ -115,7 +115,7 @@ def train(model, criterion, data_loader, optimizer, epoch):
                M[n, t] = val
        e_x = np.exp(M - np.max(M))
        M = e_x / e_x.sum(axis=0) # only difference
-        M = Variable(torch.FloatTensor(M))
+        M = Variable(torch.FloatTensor(M)).cuda()
        M = torch.stack([M]*32)

        # forward pass