fix VI: return of the fix

mosaicml · Apr 19, 2024 · b7fb56a · b7fb56a
1 parent a59e09c
commit b7fb56a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/llmfoundry/data/packing.py b/llmfoundry/data/packing.py
@@ -8,7 +8,6 @@
 import numpy as np
 import torch
 from composer.utils import dist
-from omegaconf import DictConfig
 from transformers import PreTrainedTokenizerBase
 
 log = logging.getLogger(__name__)

diff --git a/tests/data/test_packing.py b/tests/data/test_packing.py
@@ -107,9 +107,9 @@ def test_auto_packing(profile_packing: Mock):
     profile_packing.return_value = [(1, .9, 0), (2, .8, 0), (3, .7, .5)]
 
     packing_ratio = auto_packing_ratio(
-        dataloader_cfg=DictConfig({'dataset': {
+        dataloader_cfg={'dataset': {
             'max_seq_len': 2048
-        }}),
+        }},
         tokenizer=None,
         device_batch_size=1,
     )  # Dummy values, profiling results are already set.
@@ -134,9 +134,9 @@ def test_dist_auto_packing(profile_packing: Mock):
                                         (3, .7, .5)]  # should pick 2
 
     packing_ratio = auto_packing_ratio(
-        dataloader_cfg=DictConfig({'dataset': {
+        dataloader_cfg={'dataset': {
             'max_seq_len': 2048
-        }}),
+        }},
         tokenizer=None,
         device_batch_size=1,
     )  # Dummy values, profiling results are already set.