fix the text classification labels_file artifact. (#2751)

NVIDIA · Aug 31, 2021 · 122ee9a · 122ee9a
1 parent b2ace62
commit 122ee9a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/nlp/models/text_classification/text_classification_model.py b/nemo/collections/nlp/models/text_classification/text_classification_model.py
@@ -85,7 +85,7 @@ def __init__(self, cfg: DictConfig, trainer: Trainer = None):
 
         # register the file containing the labels into the artifacts to get stored in the '.nemo' file later
         if 'class_labels' in cfg and 'class_labels_file' in cfg.class_labels and cfg.class_labels.class_labels_file:
-            self.register_artifact('class_labels', cfg.class_labels.class_labels_file)
+            self.register_artifact('class_labels.class_labels_file', cfg.class_labels.class_labels_file)
 
     def create_loss_module(self):
         # create the loss module if it is not yet created by the training data loader