fix timing

black-forest-labs · Sep 9, 2024 · 170510f · 170510f
1 parent d013dc4
commit 170510f
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/flux/cli.py b/src/flux/cli.py
@@ -215,6 +215,7 @@ def main(
         x = unpack(x.float(), opts.height, opts.width)
         with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
             x = ae.decode(x)
+        torch.cuda.synchronize()
         t1 = time.perf_counter()
 
         fn = output_name.format(idx=idx)