Update app.py

app.py

@@ -6,6 +6,8 @@ import gradio as gr
 import sentencepiece
 
 
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
+
 # Load the model and tokenizer using transformers
 model = AutoModelForCausalLM.from_pretrained("01-ai/Yi-34B-200K", device_map="auto", torch_dtype="auto", trust_remote_code=True)
 tokenizer = YiTokenizer(vocab_file="./tokenizer.model")