Update src/accelerate/data_loader.py

Co-authored-by: Zach Mueller <[email protected]>
byi8220 · Jul 15, 2024 · 45db4b9 · 45db4b9
1 parent 8f04c1e
commit 45db4b9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
@@ -404,7 +404,7 @@ def __init__(self, dataset, use_stateful_dataloader=False, batch_sampler=None, *
         self.use_stateful_dataloader = use_stateful_dataloader
 
         if use_stateful_dataloader and not is_torchdata_stateful_dataloader_available():
-            raise ValueError("StatefulDataLoader is not available. Please install torchdata to use it.")
+            raise ImportError("StatefulDataLoader is not available. Please install torchdata to use it.")
         if use_stateful_dataloader:
             self.base_dataloader = StatefulDataLoader(dataset, batch_sampler=batch_sampler, **kwargs)
         else: