ssaroya
/

gptq_model

Inference Endpoints

Model card Files Files and versions Community

ssaroya commited on May 22, 2023

Commit

5e4faa3

·

1 Parent(s): 320d690

Update handler.py

Files changed (1) hide show

handler.py +3 -0

handler.py CHANGED Viewed

@@ -7,6 +7,9 @@ from utils import find_layers, DEV
 from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM
 import os
 class EndpointHandler:
     def __init__(self,
                  path="",

 from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM
 import os
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 class EndpointHandler:
     def __init__(self,
                  path="",