Update README.md

README.md

@@ -33,6 +33,7 @@ NOTE can be utilized to generate various summaries not only discharge summaries
 ## Usage
 
 ~~~python
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 model = AutoModelForCausalLM.from_pretrained("jinee/note", load_in_4bit=True, device_map="auto")
 tokenizer = AutoTokenizer.from_pretrained("jinee/note")
@@ -55,6 +56,7 @@ Also, never make up information that is not in the input data, and write a repor
 7. Discharge outcomes and treatment plan
 8. Overall summary of at least 500 characters in lines including the above contents
 '''
+torch.cuda.empty_cache()
 
 def generation(model, tokenizer, input_data):
     pipe = pipeline('text-generation',