Add pin_memory to DataLoader and update ImageInfo to support

2026-04-18 01:30:02 +00:00 · 2025-01-23 10:39:01 -05:00
parent e89653975d
commit c4b0bb6fce
17 changed files with 43 additions and 0 deletions
--- a/flux_train.py
+++ b/flux_train.py
@@ -397,6 +397,7 @@ def train(args):
        shuffle=True,
        collate_fn=collator,
        num_workers=n_workers,
+        pin_memory=args.pin_memory,
        persistent_workers=args.persistent_data_loader_workers,
    )