Adding None handle to LoRA.

Signed-off-by: Ming Huang <mingh@nvidia.com>
NVIDIA · Apr 9, 2024 · e5ccefb · e5ccefb
1 parent 9c14f84
commit e5ccefb
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py
@@ -651,7 +651,6 @@ def __eq__(self, other):
 
 
 def _canonicalize_lora_scope(scope):
-    scope = scope.lower()
 
     SCOPE_NONE = 'none'
     SCOPE_ALL = 'all'
@@ -662,6 +661,10 @@ def _canonicalize_lora_scope(scope):
     SCOPE_EX_OUTPUT_PROJ = 'exclude_output_proj'
     SCOPE_EX_MLP = 'exclude_mlp'
 
+    scope = SCOPE_NONE if scope is None else scope
+
+    scope = scope.lower()
+
     assert scope in [
         SCOPE_NONE, SCOPE_ALL, SCOPE_QKV_PROJ, SCOPE_OUTPUT_PROJ, SCOPE_MLP, SCOPE_EX_QKV_PROJ,
         SCOPE_EX_OUTPUT_PROJ, SCOPE_EX_MLP