algorithm/pytorch-tutorial
1
0
Fork 0
mirror of https://github.com/yunjey/pytorch-tutorial.git synced 2025-07-05 16:36:44 +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

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

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