Changed from **kwargs

Signed-off-by: Boris Fomitchev <bfomitchev@nvidia.com>
borisfom · Jul 6, 2023 · 251aec8 · 251aec8
1 parent d234e90
commit 251aec8
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 7 deletions.
diff --git a/nemo/collections/asr/models/asr_model.py b/nemo/collections/asr/models/asr_model.py
@@ -240,10 +240,10 @@ def disabled_deployment_input_names(self):
     def disabled_deployment_output_names(self):
         return self.encoder.disabled_deployment_output_names
 
-    def set_export_config(self, **kwargs):
-        if 'cache_support' in kwargs:
-            enable = bool(kwargs['cache_support'])
+    def set_export_config(self, args):
+        if 'cache_support' in args:
+            enable = bool(args['cache_support'])
             self.encoder.export_cache_support = enable
             logging.info(f"Caching support enabled: {enable}")
             self.encoder.setup_streaming_params()
-        super().set_export_config(**kwargs)
+        super().set_export_config(args)
diff --git a/nemo/core/classes/exportable.py b/nemo/core/classes/exportable.py
@@ -306,7 +306,7 @@ def list_export_subnets(self):
     def get_export_config(self):
         return getattr(self, 'export_config', {})
 
-    def set_export_config(self, **kwargs):
+    def set_export_config(self, kwargs):
         ex_config = self.get_export_config()
         ex_config.update(kwargs)
         self.export_config = ex_config
diff --git a/scripts/export.py b/scripts/export.py
@@ -140,11 +140,11 @@ def nemo_export(argv):
         max_dim = args.max_dim
 
     if args.cache_support:
-        model.set_export_config(cache_support=True)
+        model.set_export_config({cache_support: True})
 
     if args.config:
         kv = dict(map(lambda s: s.split('='), args.config))
-        model.set_export_config(**kv)
+        model.set_export_config(kv)
 
     autocast = nullcontext
     if args.autocast: