fix in ltx2_pipeline

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -92,7 +92,7 @@ def cast_with_exclusion(path, x, dtype_to_cast):
     return x.astype(dtype_to_cast)
 
 
-def _add_sharding_rule(vs: nnx.VariableState, logical_axis_rules) -> nnx.VariableState:
+def _add_sharding_rule(vs: nnx.Variable, logical_axis_rules) -> nnx.Variable:
   vs.sharding_rules = logical_axis_rules
   return vs