From 536da66dadab4b54e9017fb2b4d09b2b916fe747 Mon Sep 17 00:00:00 2001 From: Pavel Belevich Date: Tue, 8 Sep 2020 17:16:24 -0400 Subject: [PATCH] Make FCModel.py working on CUDA --- models/FCModel.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/models/FCModel.py b/models/FCModel.py index c275b5b9..604b1140 100644 --- a/models/FCModel.py +++ b/models/FCModel.py @@ -59,10 +59,10 @@ def __init__(self, opt): self.ss_prob = 0.0 # Schedule sampling probability - self.img_embed = nn.Linear(self.fc_feat_size, self.input_encoding_size) - self.core = LSTMCore(opt) - self.embed = nn.Embedding(self.vocab_size + 1, self.input_encoding_size) - self.logit = nn.Linear(self.rnn_size, self.vocab_size + 1) + self.img_embed = nn.Linear(self.fc_feat_size, self.input_encoding_size).to(device) + self.core = LSTMCore(opt).to(device) + self.embed = nn.Embedding(self.vocab_size + 1, self.input_encoding_size).to(device) + self.logit = nn.Linear(self.rnn_size, self.vocab_size + 1).to(device) self.init_weights()