From f1878cb524f8fbc844e06776991fc498c649a11e Mon Sep 17 00:00:00 2001 From: "majiahui@haimaqingfan.com" Date: Wed, 12 Jul 2023 18:30:01 +0800 Subject: [PATCH] =?UTF-8?q?=E8=A1=A5=E5=85=85readme?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- flask_batch.py | 2 ++ predict.py | 1 - 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/flask_batch.py b/flask_batch.py index fbd8be2..16f9e20 100644 --- a/flask_batch.py +++ b/flask_batch.py @@ -1,3 +1,5 @@ +import os +os.environ["CUDA_VISIBLE_DEVICES"] = "0" from flask import Flask, jsonify from flask import request from transformers import pipeline diff --git a/predict.py b/predict.py index 4ac8318..799db26 100644 --- a/predict.py +++ b/predict.py @@ -1,5 +1,4 @@ import time - from vllm import LLM, SamplingParams prompts = [