fix use cuda

convlab/modules/nlg/multiwoz/sc_lstm/model/layers/decoder_deep.py

@@ -26,17 +26,30 @@ def __init__(self, dec_type, input_size, output_size, hidden_size, d_size, n_lay
  self.w2h, self.h2h = nn.ModuleList(), nn.ModuleList()
  self.w2h_r, self.h2h_r = nn.ModuleList(), nn.ModuleList()
  self.dc = nn.ModuleList()
- for i in range(n_layer):
- if i == 0:
- self.w2h.append( nn.Linear(input_size, hidden_size*4).cuda() )
- self.w2h_r.append( nn.Linear(input_size, d_size).cuda() )
- else:
- self.w2h.append( nn.Linear(input_size + i*hidden_size, hidden_size*4).cuda() )
- self.w2h_r.append( nn.Linear(input_size + i*hidden_size, d_size).cuda() )
-
- self.h2h.append( nn.Linear(hidden_size, hidden_size*4).cuda() )
- self.h2h_r.append( nn.Linear(hidden_size, d_size).cuda() )
- self.dc.append( nn.Linear(d_size, hidden_size, bias=False).cuda() )
+ if use_cuda:
+ for i in range(n_layer):
+ if i == 0:
+ self.w2h.append( nn.Linear(input_size, hidden_size*4).cuda() )
+ self.w2h_r.append( nn.Linear(input_size, d_size).cuda() )
+ else:
+ self.w2h.append( nn.Linear(input_size + i*hidden_size, hidden_size*4).cuda() )
+ self.w2h_r.append( nn.Linear(input_size + i*hidden_size, d_size).cuda() )
+
+ self.h2h.append( nn.Linear(hidden_size, hidden_size*4).cuda() )
+ self.h2h_r.append( nn.Linear(hidden_size, d_size).cuda() )
+ self.dc.append( nn.Linear(d_size, hidden_size, bias=False).cuda() )
+ else:
+ for i in range(n_layer):
+ if i == 0:
+ self.w2h.append( nn.Linear(input_size, hidden_size*4) )
+ self.w2h_r.append( nn.Linear(input_size, d_size) )
+ else:
+ self.w2h.append( nn.Linear(input_size + i*hidden_size, hidden_size*4) )
+ self.w2h_r.append( nn.Linear(input_size + i*hidden_size, d_size) )
+
+ self.h2h.append( nn.Linear(hidden_size, hidden_size*4) )
+ self.h2h_r.append( nn.Linear(hidden_size, d_size) )
+ self.dc.append( nn.Linear(d_size, hidden_size, bias=False) )
 
  self.out = nn.Linear(hidden_size*n_layer, output_size)
 

convlab/modules/nlg/multiwoz/sc_lstm/nlg_sc_lstm.py

@@ -60,6 +60,7 @@ def __init__(self,
  archive = zipfile.ZipFile(archive_file, 'r')
  archive.extractall(model_dir)
 
+ self.USE_CUDA = use_cuda
  self.args, self.config = parse(is_user)
  self.dataset = SimpleDatasetWoz(self.config)
 
@@ -76,7 +77,8 @@ def __init__(self,
  assert os.path.isfile(model_path)
  self.model.load_state_dict(torch.load(model_path))
  self.model.eval()
- self.model.cuda()
+ if use_cuda:
+ self.model.cuda()
 
  def generate_delex(self, meta):
  """
@@ -126,7 +128,8 @@ def generate_delex(self, meta):
  feats = [do_cond + da_cond + sv_cond]
 
  feats_var = torch.FloatTensor(feats)
- feats_var = feats_var.cuda()
+ if self.USE_CUDA:
+ feats_var = feats_var.cuda()
 
  decoded_words = self.model.generate(self.dataset, feats_var, self.args['beam_size'])
  delex = decoded_words[0] # (beam_size)