From 2fe796bb100081db3d61d5e42711947f358a4f7a Mon Sep 17 00:00:00 2001 From: yunjey Date: Wed, 22 Mar 2017 22:58:29 +0900 Subject: [PATCH] rearrange code for cuda() --- tutorials/09 - Image Captioning/train.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tutorials/09 - Image Captioning/train.py b/tutorials/09 - Image Captioning/train.py index eeb41e5..d167fd1 100644 --- a/tutorials/09 - Image Captioning/train.py +++ b/tutorials/09 - Image Captioning/train.py @@ -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()