Update backend_model.py

backend_model.py

@@ -110,9 +110,8 @@ def load_model_and_tokenizer(model_architecture, model_path):
     model = CodeBERTBinaryClassifier(base_model)
     # model = model.to(device)
 
-    # Load the model
-    # model = CodeBERTBinaryClassifier(base_model)
-    model.load_state_dict(torch.load(model_path))
-    model = model.to(device)
+    map_location = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model.load_state_dict(torch.load(model_path, map_location=map_location))
+    model = model.to(map_location)
 
     return model, tokenizer