Fix bug in get_update_constraint (#125194)

Summary: Title

Test Plan: CI

Differential Revision: D56726321

Pull Request resolved: #125194
Approved by: https://github.com/pianpwk

torch/export/exported_program.py

@@ -823,15 +823,15 @@ def get_shape_env(gm):
 
  fake_mode = detect_fake_mode(vals)
  if fake_mode is not None:
- return fake_mode.shape_env, fake_mode
+ return fake_mode.shape_env
  for v in vals:
  if isinstance(v, torch.SymInt):
- return v.node.shape_env, fake_mode
+ return v.node.shape_env
 
  # FIXME(tmanlaibaatar) Remove this whole branch once https://github.com/pytorch/pytorch/pull/123764
  if _is_executorch:
  assert old_range_constraints is None
- shape_env, _ = get_shape_env(gm)
+ shape_env = get_shape_env(gm)
  if shape_env is None:
  return {}
  range_constraints = {
@@ -849,7 +849,7 @@ def get_shape_env(gm):
 
  assert old_range_constraints is not None
 
- shape_env, fake_mode = get_shape_env(gm)
+ shape_env = get_shape_env(gm)
  if shape_env is None:
  return {}