debugging for cuda

new_test_saved_finetuned_model.py

@@ -55,7 +55,7 @@ class BERTFineTuneTrainer:
         # available_gpus = list(range(torch.cuda.device_count()))
 
         # This BERT model will be saved every epoch
-        self.model = bertFinetunedClassifierwithFeats.to("cpu")
+        self.model = bertFinetunedClassifierwithFeats.to(self.device)
         print(self.model.parameters())
         for param in self.model.parameters():
             param.requires_grad = False
@@ -159,7 +159,7 @@ class BERTFineTuneTrainer:
                     logits = self.model.forward(data["input"], data["segment_label"], data["feat"])
                 else:
                     with torch.no_grad():
-                        logits = self.model.forward(data["input"].cpu(), data["segment_label"].cpu(), data["feat"].cpu())
+                        logits = self.model.forward(data["input"], data["segment_label"], data["feat"])
 
                 logits = logits.cpu()
                 loss = self.criterion(logits, data["label"])