layerfreezing

mosaicml · Nov 30, 2024 · 53da8a8 · 53da8a8
1 parent a46bcc5
commit 53da8a8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_full_nlp.py b/tests/test_full_nlp.py
@@ -11,7 +11,7 @@
 from torchmetrics.classification import MulticlassAccuracy
 from transformers import BertConfig, BertForMaskedLM, BertForSequenceClassification, BertTokenizerFast
 
-from composer.algorithms import LabelSmoothing
+from composer.algorithms import LayerFreezing
 from composer.loggers import RemoteUploaderDownloader
 from composer.metrics.nlp import LanguageCrossEntropy, MaskedAccuracy
 from composer.models import HuggingFaceModel
@@ -233,7 +233,7 @@ def inference_test_helper(
 @pytest.mark.parametrize(
     'model_type,algorithms,save_format',
     [
-        ('tinybert_hf', [LabelSmoothing], 'onnx'),
+        ('tinybert_hf', [LayerFreezing], 'onnx'),
         ('simpletransformer', [], 'torchscript'),
     ],
 )