From 4645e28355363d5dceb5f644a0c1ea5bdc2471f9 Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Mon, 12 Dec 2022 04:56:02 -0800 Subject: [PATCH] tensor format ort bug fix (#1557) bug fix Co-authored-by: Prathik Rao Co-authored-by: anton- --- .../unconditional_image_generation/train_unconditional_ort.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/unconditional_image_generation/train_unconditional_ort.py b/examples/unconditional_image_generation/train_unconditional_ort.py index 71b192f49a..4e97732ade 100644 --- 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,