From 2e295c9f94dbefb81572b3d9ecf68fe014cb9d2a Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Wed, 19 Jul 2023 22:58:16 -0400
Subject: [PATCH] use accelerator prepare for dataloader

---
 src/axolotl/utils/trainer.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/axolotl/utils/trainer.py b/src/axolotl/utils/trainer.py
index cc68edb02..5009d645b 100644
--- a/src/axolotl/utils/trainer.py
+++ b/src/axolotl/utils/trainer.py
@@ -101,12 +101,14 @@ class AxolotlTrainer(Trainer):
         if self.args.sample_packing:
             train_sampler = self._get_train_sampler()
 
-            return MultipackDistributedDataloader(
-                self.train_dataset,
-                batch_size=self._train_batch_size,
-                seq_max_length=self.args.max_seq_length,
-                collate_fn=self.data_collator,
-                sampler=train_sampler,
+            return self.accelerator.prepare(
+                MultipackDistributedDataloader(
+                    self.train_dataset,
+                    batch_size=self._train_batch_size,
+                    seq_max_length=self.args.max_seq_length,
+                    collate_fn=self.data_collator,
+                    sampler=train_sampler,
+                )
             )
         return super().get_train_dataloader()