Repository: pytorch
Issue: LTC asynchronous executor is not synced with python code and ignores context manager.
Labels: triaged,LazyTensor_nvfuser_integration,module: lazy

## 🐛 Bug

I'm observing a behavior where LTC is not respecting python context manager.

e.g.
```
# enabling nvfuser instead of NNC for TorchScript fusion
with torch.jit.fuser("fuser2"):
  for i in range(ITER):
    # running lazy tensor
    ...
    ltm.mark_step()
```

I noticed that in the executor that's executing the lazy tensor computation is using NNC fuser instead of nvfuser. Basically we are ignoring the `with torch.jit.fuser("fuser2")` context manager that enabled nvfuser and disabled NNC fuser.

My simple printf debugging shows that we are exiting the `with.jit.fuser` before executing the TorchScript. So by the time that we run TorchScript executor, we've switched back to NNC fuser already.

## To Reproduce

Running the script with vvv
`LTC_TS_CUDA=1 PYTORCH_NVFUSER_DISABLE_FALLBACK=1 PYTORCH_JIT_LOG_LEVEL="profiling_graph_executor_impl" python nvfuser_example.py &> cache`

Here's nvfuser_example.py
```
import torch                                                                                                               
import lazy_tensor_core                                                                                                    
import lazy_tensor_core.debug.metrics as metrics                                                                           
import lazy_tensor_core.core.lazy_model as ltm                                                                             
                                                                                                                           
lazy_tensor_core._LAZYC._ltc_init_ts_backend()                                                                             
                                                                                                                           
# using this instead of the with statement does resolve the problem.                                                       
# torch._C._jit_set_nvfuser_enabled(True)                                                                                  
# torch._C._jit_set_texpr_fuser_enabled(False)                                                                             
# torch._C._jit_set_profiling_executor(True)                                                                               
# torch._C._jit_set_profiling_mode(True)                                                                                   
# torch._C._jit_override_can_fuse_on_cpu(False)                                                                            
# torch._C._jit_override_can_fuse_on_gpu(False)                                                                            
# torch._C._jit_set_bailout_depth(20)                                                                                      
                                                                                                                           
torch.manual_seed(42)                                                                                                      
                                                                                                                           
device = 'lazy'                                                                                                            
dtype = torch.float32                                                                                                      
ITER = 5                                                                                                                   
                                                                                                                           
x = torch.randn(2, 4, device=device, dtype=dtype)                                                                          
                                                                                                                           
class MyModule(torch.nn.Module):                                                                                           
  def __init__(self):                                                                                                      
    super(MyModule, self).__init__()                                                                                       
    self.linear = torch.nn.Linear(4, 4)                                                                                    
    self.relu = torch.nn.ReLU()                                                                                            
  def forward(self, x):                                                                                                    
    o = self.linear(x)                                                                                                     
    o = self.relu(o)                                                                                                       
    return o                                                                                                               
                                                                                                                           
module = MyModule().to(device=device, dtype=dtype)                                                                         
optim = torch.optim.SGD(module.parameters(), lr = 1e-2, momentum = 0.0)                                                    
loss = torch.nn.CrossEntropyLoss()                                                                                         
                                                                                                                           
# looks like by the time we are executing the TorchScript, we've already exited                                            
# this context manager and default fuser back to NNC.                                                                                                                                                              
with torch.jit.fuser("fuser2"):                                                                                            
  for i in range(ITER):                                                                                                    
    o = module(x)                                                                                                          
    target = torch.empty(2, dtype=torch.long, device=device).random_(4)                                                    
    loss(o, target).backward()                                                                                             
    optim.step()                                                                                                           
    ltm.mark_step()                                                                                                        
  print(metrics.metrics_report())
```

You'll notice that in the cache file, we are seeing `prim::TensorExprGroup` node in the graph, although we've disabled that and switched to nvfuser instead.
If you switch from the context manager back to torch._C APIs (uncomment the torch._C... and comment out the with statement), you'll observe that we are not using `prim::CudaFusionGroup`

## Environment

I'm on lazy_tensor_staging branch
```
commit 43b71a1a5ec24924e4c7bedd573865e74b0411a4 (HEAD, upstream/lazy_tensor_staging)
Author: Bin Bao <binbao@fb.com>
Date:   Thu Oct 28 18:58:09 2021 +0000

    Remove literal_util.h
```

## Additional context

I've tried the sync flag `SYNC_TENSORS_OPBYOP` trying to avoid the async execution, hoping that would save me from the issue, but that runs into a segfault with error:
```
root@b1a4ac0e590f:/raid/playground/LTC# LTC_TS_CUDA=1 PYTORCH_NVFUSER_DISABLE_FALLBACK=1 SYNC_TENSORS_OPBYOP=1 PYTORCH_JIT_LOG_LEVEL="profiling_graph_executor_impl" python nvfuser_example.py &> cache
Aborted (core dumped)
root@b1a4ac0e590f:/raid/playground/LTC# vim cache

[F lowering_context.cpp:33] Not implemented.
```