Artix

Runtime error

Vitrous commited on Feb 15, 2024

Commit

d2c9447

verified ·

1 Parent(s): 4f9ba28

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -1,6 +1,8 @@
 import uvicorn
 from fastapi import FastAPI, HTTPException, Request
 from auto_gptq import AutoGPTQForCausalLM
 import torch
 import optimum
 from transformers import (AutoModelForCausalLM, AutoTokenizer, LlamaForCausalLM, LlamaTokenizer, GenerationConfig, pipeline,)

 import uvicorn
 from fastapi import FastAPI, HTTPException, Request
 from auto_gptq import AutoGPTQForCausalLM
+import os
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
 import torch
 import optimum
 from transformers import (AutoModelForCausalLM, AutoTokenizer, LlamaForCausalLM, LlamaTokenizer, GenerationConfig, pipeline,)