From 74809cec3799c5b72ab912da937e74e6bf234ae8 Mon Sep 17 00:00:00 2001
From: Xun Liu <30335480+AntiQuality@users.noreply.github.com>
Date: Thu, 7 Mar 2024 14:58:58 +0800
Subject: [PATCH] Complete missing `import`

---
 README.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/README.md b/README.md
index 321c9d0..77ad8f7 100644
--- a/README.md
+++ b/README.md
@@ -129,6 +129,7 @@ This code will output the following result:
 #### 3) Chat Model Inference
 ```python
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-coder-6.7b-instruct", trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-coder-6.7b-instruct", trust_remote_code=True, torch_dtype=torch.bfloat16).cuda()
 messages=[
@@ -177,6 +178,7 @@ You are an AI programming assistant, utilizing the DeepSeek Coder model, develop
 #### 4) Repository Level Code Completion
 ```python
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-coder-6.7b-base", trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-coder-6.7b-base", trust_remote_code=True, torch_dtype=torch.bfloat16).cuda()