Fixed minor bug when running training on cuda

examples/run_pplm_discrim_train.py

@@ -18,6 +18,7 @@ import torch.utils.data as data
 from nltk.tokenize.treebank import TreebankWordDetokenizer
 from torchtext import data as torchtext_data
 from torchtext import datasets
+
 from transformers import GPT2Tokenizer, GPT2LMHeadModel
 
 torch.manual_seed(0)
@@ -89,7 +90,7 @@ class Discriminator(torch.nn.Module):
         if self.cached_mode:
             avg_hidden = x.to(device)
         else:
-            avg_hidden = self.avg_representation(x)
+            avg_hidden = self.avg_representation(x.to(device))
 
         logits = self.classifier_head(avg_hidden)
         probs = F.log_softmax(logits, dim=-1)
@@ -203,7 +204,7 @@ def evaluate_performance(data_loader, discriminator):
 
 def predict(input_sentence, model, classes, cached=False):
     input_t = model.tokenizer.encode(input_sentence)
-    input_t = torch.tensor([input_t], dtype=torch.long)
+    input_t = torch.tensor([input_t], dtype=torch.long, device=device)
     if cached:
         input_t = model.avg_representation(input_t)
 
@@ -428,6 +429,7 @@ def train_discriminator(
         with open(dataset_fp) as f:
             csv_reader = csv.reader(f, delimiter='\t')
             for row in csv_reader:
+                if row:
                     classes.add(row[0])
 
         idx2class = sorted(classes)
@@ -444,6 +446,7 @@ def train_discriminator(
         with open(dataset_fp) as f:
             csv_reader = csv.reader(f, delimiter='\t')
             for i, row in enumerate(csv_reader):
+                if row:
                     label = row[0]
                     text = row[1]