fix lint

orttraining/orttraining/python/training/ort_triton/_codegen.py

@@ -456,7 +456,7 @@ def ModuleNode(self, node: ModuleNode, context: CodegenContext, code_buffer: Cod
         for input in node.inputs:
             for idx, dim in enumerate(input.shape):
                 if dim.is_symbol and dim not in seen_symbolic_shape:
-                    code_buffer += f"{space_indent}{str(dim)} = {context.get_variable_name(input.name)}.size()[{idx}]\n"
+                    code_buffer += f"{space_indent}{dim} = {context.get_variable_name(input.name)}.size()[{idx}]\n"
                     seen_symbolic_shape.add(dim)
 
         if node.has_dropout:

orttraining/orttraining/python/training/ort_triton/_common.py

@@ -12,14 +12,14 @@
 from ._sympy_utils import extract_shape_from_symbol
 from ._utils import get_attribute, get_reduce_info, next_power_of_2
 
+_SPECIAL_FLOATS: List[str] = ["inf", "-inf"]
+
 
 class CodegenContext:
     """
     record variable name mapping in term of IRnodes.
     """
 
-    _special_floats = ["inf", "-inf"]
-
     def __init__(self, var_map: Dict[str, str]):
         self._var_map: Dict[str, str] = {**var_map}
 
@@ -31,7 +31,7 @@ def get_variable_name(self, name: str) -> str:
     def get_internal_variable_name(self, name: str) -> str:
         var_name = self._var_map[name]
         var_name = self._var_map[var_name] if var_name in self._var_map else var_name
-        return f'float("{var_name}")' if var_name in self._special_floats else var_name
+        return f'float("{var_name}")' if var_name in _SPECIAL_FLOATS else var_name
 
 
 class CodeBuffer: