Hi @kyong,
FX does not support context managers (such as no_grad) in traced regions. Can you refactor the original code to something like:
class MyModel(torch.nn.Module):
  def forward(self, x):
    tmp_val = … do something …
    out = torch.mul(x, tmp_val)
    return out

class NoGradWrapper(torch.nn.Module):
  def __init__(self):
    self.mm = MyModel()

  def forward(self, x):
    with torch.no_grad():
      return self.mm(x)

Then you can apply FX transforms to NoGradWrapper.mm and still preserve the no-gradient behavior