Update main.py

main.py

@@ -13,6 +13,11 @@ name = "microsoft/DialoGPT-medium"
 model = GPT2LMHeadModel.from_pretrained(name)
 tokenizer = GPT2Tokenizer.from_pretrained(name)
 
+# Using CUDA for an optimal experience
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = model.to(device)
+
+# Open a thing for the API
 @app.route("/api", methods=["POST"])
 def receive_data():
   data = request.get_json()
@@ -33,7 +38,8 @@ def receive_data():
   print("Answered with:", answer_data)
   return jsonify(answer_data)
 
-@app.route("/", methods=["GET"])
+# Incase a normal browser opens the page
+@app.route("/")
 def not_api():
   return render_template("index.html")