tensor format ort bug fix (#1557)

bug fix Co-authored-by: Prathik Rao <prathikrao@microsoft.com> Co-authored-by: anton- <anton@huggingface.co>

tensor format ort bug fix (#1557)
bug fix Co-authored-by: Prathik Rao <prathikrao@microsoft.com> Co-authored-by: anton- <anton@huggingface.co>
4645e283 · Prathik Rao · GitHub · 58933059 · 4645e283
Unverified Commit 4645e283 authored Dec 12, 2022 by Prathik Rao Committed by GitHub Dec 12, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

examples/unconditional_image_generation/train_unconditional_ort.py ...unconditional_image_generation/train_unconditional_ort.py +1 -1

No files found.
--- a/examples/unconditional_image_generation/train_unconditional_ort.py
+++ b/examples/unconditional_image_generation/train_unconditional_ort.py
@@ -77,7 +77,7 @@ def main(args):
        ),
    )
    model = ORTModule(model)
-    noise_scheduler = DDPMScheduler(num_train_timesteps=1000, tensor_format="pt")
+    noise_scheduler = DDPMScheduler(num_train_timesteps=1000)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=args.learning_rate,