WenhaoWang commited on
Commit
62ed335
1 Parent(s): 5c8879d

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +2 -1
README.md CHANGED
@@ -22,7 +22,8 @@ It is fine-tuned on the [VidProM](https://huggingface.co/datasets/WenhaoWang/Vid
22
  ## Download the model
23
  ```
24
  from transformers import pipeline
25
- pipe = pipeline("text-generation", model="WenhaoWang/AutoT2VPrompt")
 
26
  ```
27
 
28
  ## Set the Parameters
 
22
  ## Download the model
23
  ```
24
  from transformers import pipeline
25
+ import torch
26
+ pipe = pipeline("text-generation", model="WenhaoWang/AutoT2VPrompt", model_kwargs={"torch_dtype": torch.bfloat16}, device_map="cuda:0")
27
  ```
28
 
29
  ## Set the Parameters