old_prof_exec_state = torch._C._jit_set_profiling_executor(True)
old_prof_mode_state = torch._C._jit_set_profiling_mode(True)

this enables the profiling executor.

torch._C._jit_set_num_profiled_runs(num_runs)

how many profiling runs we want to do before we optimize the graph.
@torch.jit.script
def foo(x):
    if x.size(0) == 1:
        return 1
    else:
        return 2

old_prof_exec_state = torch._C._jit_set_profiling_executor(True)
old_prof_mode_state = torch._C._jit_set_profiling_mode(True)
torch._C._jit_set_num_profiled_runs(1)

foo(torch.rand([1, 2]))
print(torch.jit.last_executed_optimized_graph())
foo(torch.rand([1, 2]))
print(torch.jit.last_executed_optimized_graph())

gives
graph(%x.1 : Tensor):
  %1 : int = prim::Constant[value=0]() # test/test_jit.py:3742:22
  %2 : int = prim::Constant[value=1]() # test/test_jit.py:3742:28
  %3 : int = prim::Constant[value=2]() # test/test_jit.py:3745:23
  %4 : Tensor = prim::profile(%x.1)
  %5 : int = aten::size(%4, %1) # test/test_jit.py:3742:15
  %6 : bool = aten::eq(%5, %2) # test/test_jit.py:3742:15
  %7 : int = prim::If(%6) # test/test_jit.py:3742:12
    block0():
      -> (%2)
    block1():
      -> (%3)
   = prim::profile()
  return (%7)

first graph, still profiling
graph(%x.1 : Tensor):
  %2 : int = prim::Constant[value=1]() # test/test_jit.py:3742:28
  %10 : int = prim::BailoutTemplate_0()
  %9 : Double(1:2, 2:1) = prim::BailOut[index=0](%10, %x.1)
  return (%2)
with prim::BailoutTemplate_0 = graph(%x.1 : Tensor):
  %1 : Double(1:2, 2:1) = prim::BailOut[index=0](%x.1)
  %2 : int = prim::Constant[value=0]() # test/test_jit.py:3742:22
  %3 : int = prim::Constant[value=1]() # test/test_jit.py:3742:28
  %4 : int = prim::Constant[value=2]() # test/test_jit.py:3745:23
  %5 : int = aten::size(%1, %2) # test/test_jit.py:3742:15
  %6 : bool = aten::eq(%5, %3) # test/test_jit.py:3742:15
  %7 : int = prim::If(%6) # test/test_jit.py:3742:12
    block0():
      -> (%3)
    block1():
      -> (%4)
  return (%7)

second graph, optimized from profiles.
additionally, you can read
https://github.com/pytorch/pytorch/blob/53af9df557aff745edf24193ece784fd008c6f19/torch/csrc/jit/OVERVIEW.md#profiling-programs 16.