Skip to content

Commit fe0992e

Browse files
committed
Add Torch CUDA sync to fix timing code in cli.py
1 parent d013dc4 commit fe0992e

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

src/flux/cli.py

+3
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,9 @@ def main(
215215
x = unpack(x.float(), opts.height, opts.width)
216216
with torch.autocast(device_type=torch_device.type, dtype=torch.bfloat16):
217217
x = ae.decode(x)
218+
219+
if torch.cuda.is_available():
220+
torch.cuda.synchronize()
218221
t1 = time.perf_counter()
219222

220223
fn = output_name.format(idx=idx)

0 commit comments

Comments
 (0)