Move everything to CPU

ModelDriver.py

@@ -5,7 +5,8 @@ import torch.nn.functional as F
 from torch.utils.data import TensorDataset, DataLoader
 
 
-device = torch.device("cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# device = torch.device("cpu")
 class MLP(nn.Module):
     def __init__(self, input_dim):
         super(MLP, self).__init__()
@@ -62,14 +63,14 @@ def RobertaClassifierOpenGPTInference(input_text):
     tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
     model_path = "ClassifierCheckpoint/RobertaClassifierOpenGPT.pth"
     model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=2)
-    model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
-    model = model.to(torch.device('cpu'))
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model = model.to(device)
     model.eval()
 
 
     tokenized_input = tokenizer(input_text, truncation=True, padding=True, max_length=512, return_tensors='pt')
-    input_ids = tokenized_input['input_ids'].to(torch.device('cpu'))
-    attention_mask = tokenized_input['attention_mask'].to(torch.device('cpu'))
+    input_ids = tokenized_input['input_ids'].to(device)
+    attention_mask = tokenized_input['attention_mask'].to(device)
 
     # Make a prediction
     with torch.no_grad():
@@ -84,14 +85,14 @@ def RobertaClassifierCSAbstractInference(input_text):
     tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
     model_path = "ClassifierCheckpoint/RobertaClassifierCSAbstract.pth"
     model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=2)
-    model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
-    model = model.to(torch.device('cpu'))
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model = model.to(device)
     model.eval()
 
 
     tokenized_input = tokenizer(input_text, truncation=True, padding=True, max_length=512, return_tensors='pt')
-    input_ids = tokenized_input['input_ids'].to(torch.device('cpu'))
-    attention_mask = tokenized_input['attention_mask'].to(torch.device('cpu'))
+    input_ids = tokenized_input['input_ids'].to(device)
+    attention_mask = tokenized_input['attention_mask'].to(device)
 
     # Make a prediction
     with torch.no_grad():

Test.py

@@ -20,7 +20,7 @@ Input_Text = "I want to do this data"
 # print(f"Confidence:", max(Probs))
 
 print("RobertaClassifierCSAbstractInference")
-Probs = RobertaClassifierOpenGPTInference(Input_Text)
+Probs = RobertaClassifierCSAbstractInference(Input_Text)
 Pred = "Human Written" if not np.argmax(Probs) else "Machine Generated"
 
 print(Probs)