Move everything to CPU 2

ModelDriver.py

@@ -61,8 +61,8 @@ def RobertaSentinelCSAbstractInference(input_text):
 def RobertaClassifierOpenGPTInference(input_text):
     tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
     model_path = "ClassifierCheckpoint/RobertaClassifierOpenGPT.pth"
-    model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=2, map_location=device)
-    model.load_state_dict(torch.load(model_path))
+    model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=2)
+    model.load_state_dict(torch.load(model_path), map_location=device)
     model = model.to(torch.device('cpu'))
     model.eval()
 
@@ -83,8 +83,8 @@ def RobertaClassifierOpenGPTInference(input_text):
 def RobertaClassifierCSAbstractInference(input_text):
     tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
     model_path = "ClassifierCheckpoint/RobertaClassifierCSAbstract.pth"
-    model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=2, map_location=device)
-    model.load_state_dict(torch.load(model_path))
+    model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=2)
+    model.load_state_dict(torch.load(model_path, map_location=device))
     model = model.to(torch.device('cpu'))
     model.eval()