From 8a3e8bf0a90d929e70b8f7e75e27e92ffbd406c8 Mon Sep 17 00:00:00 2001 From: 1000960000 user Date: Tue, 6 Aug 2024 08:58:39 +0000 Subject: [PATCH] fixes to data-collator patching for padding-free plugin Signed-off-by: 1000960000 user --- .../src/fms_acceleration_ilab/framework_plugin_padding_free.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/instruct-lab/src/fms_acceleration_ilab/framework_plugin_padding_free.py b/plugins/instruct-lab/src/fms_acceleration_ilab/framework_plugin_padding_free.py index c5725959..33a592ee 100644 --- a/plugins/instruct-lab/src/fms_acceleration_ilab/framework_plugin_padding_free.py +++ b/plugins/instruct-lab/src/fms_acceleration_ilab/framework_plugin_padding_free.py @@ -207,7 +207,7 @@ def prepare(self, *args, device_placement=None): # Replace the collate_fn in dataloader dataloader.collate_fn = DataCollatorWithFlattening() - return dataloader + return _old_prepare(dataloader) accelerator.prepare = MethodType(prepare, accelerator)