algorithm/pytorch-tutorial
1
0
Fork 0
mirror of https://github.com/yunjey/pytorch-tutorial.git synced 2025-07-06 01:15:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

rearrange code for cuda()

Browse Source
This commit is contained in:
yunjey
2017-03-22 22:58:29 +09:00
parent 5061df5123
commit 2fe796bb10
1 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tutorials/09 - Image Captioning/train.py
View File

@ -55,11 +55,10 @@ def main():
# Set mini-batch dataset # Set mini-batch dataset
images = Variable(images) images = Variable(images)
captions = Variable(captions) captions = Variable(captions)
targets = pack_padded_sequence(captions, lengths, batch_first=True)[0]
if torch.cuda.is_available(): if torch.cuda.is_available():
images = images.cuda() images = images.cuda()
captions = captions.cuda() captions = captions.cuda()
targets = pack_padded_sequence(captions, lengths, batch_first=True)[0]
# Forward, Backward and Optimize # Forward, Backward and Optimize
decoder.zero_grad() decoder.zero_grad()