I tried to run example.py on an A100 (80GB) GPU. It seems there is a bug at line [41]
|
input_ids = tokenizer(prompt, return_tensors="pt").input_ids |
The current implementation doesn't load the input_ids tensors onto the device, which causes an error. I replaced the above code, and it's now working. Fixed the issue by adding: input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")