Memory optimizations. Default to using cudamalloc when torch 2.0 for mem allocation

2023-09-12 04:30:23 -06:00
parent e8583860ad
commit d74dd636ee
5 changed files with 104 additions and 5 deletions
--- a/toolkit/data_loader.py
+++ b/toolkit/data_loader.py
@@ -484,14 +484,14 @@ def get_dataloader_from_datasets(
            drop_last=False,
            shuffle=True,
            collate_fn=dto_collation,  # Use the custom collate function
-            num_workers=1
+            num_workers=0
        )
    else:
        data_loader = DataLoader(
            concatenated_dataset,
            batch_size=batch_size,
            shuffle=True,
-            num_workers=1,
+            num_workers=0,
            collate_fn=dto_collation
        )
    return data_loader