diff --git a/web_demo.py b/web_demo.py
index 71d3257..922df00 100755
--- a/web_demo.py
+++ b/web_demo.py
@@ -9,7 +9,8 @@ from argparse import ArgumentParser
 
 import gradio as gr
 import mdtex2html
-import torch.cuda
+
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers.generation import GenerationConfig