Muhammad Taqi Raza
commited on
Commit
·
bd4af33
1
Parent(s):
ea8870e
add app files
Browse files- inference_script.py +1 -1
inference_script.py
CHANGED
|
@@ -587,7 +587,7 @@ if __name__ == "__main__":
|
|
| 587 |
# add device_map="balanced" in the from_pretrained function and remove the enable_model_cpu_offload()
|
| 588 |
# function to use Multi GPUs.
|
| 589 |
|
| 590 |
-
pipe = CogVideoXPipeline.from_pretrained(args.model_path, torch_dtype=dtype)
|
| 591 |
|
| 592 |
# If you're using with lora, add this code
|
| 593 |
if args.lora_path:
|
|
|
|
| 587 |
# add device_map="balanced" in the from_pretrained function and remove the enable_model_cpu_offload()
|
| 588 |
# function to use Multi GPUs.
|
| 589 |
|
| 590 |
+
pipe = CogVideoXPipeline.from_pretrained(args.model_path, torch_dtype=dtype, device_map="balanced")
|
| 591 |
|
| 592 |
# If you're using with lora, add this code
|
| 593 |
if args.lora_path:
|