(self, gpu_num=1)
| 8 | # 按照自己推理需求自己修改代码 |
| 9 | |
| 10 | def __init__(self, gpu_num=1): |
| 11 | model = AutoModelForCausalLM.from_pretrained(model_name) |
| 12 | device_list = [] |
| 13 | for gpu_idx in range(gpu_num): |
| 14 | device_list.append(torch.device("cuda:0")) |
| 15 | |
| 16 | # 将模型移动到指定的GPU设备 |
| 17 | model.to(device) |
| 18 | |
| 19 | |
| 20 | def gen_answer(self, chat_dict, gpu_index): |
nothing calls this directly
no outgoing calls
no test coverage detected