[fake tensor] Fake tensor is leaked out with new with Aot Autograd

[anijain2305]

Repro

import torch
from torch.nn import *

import torchdynamo


class Bar(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, argsort: torch.Tensor):
        new = argsort.new(2, 12, 4096)
        x = torch.add(new, 2)
        return (
            new,
            x,
        )


mod = Bar().to(device="cpu")

inp = torch.randn((2, 12, 4096), device="cpu")


def fn(x):
    y = mod(x)
    print("break", y[0].shape)
    return [torch.sin(l) for l in y]


opt_mod = torchdynamo.optimize("aot_eager")(fn)
opt_mod(inp)

Error

  File "/scratch/anijain/work/torchdynamo/torchdynamo/output_graph.py", line 387, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/scratch/anijain/work/torchdynamo/torchdynamo/output_graph.py", line 420, in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e) from e
torchdynamo.exc.BackendCompilerFailed: compile_fn raised Exception: Invoking operators with non-Fake Tensor inputs in FakeTensorMode is not yet supported. Please convert all Tensors to FakeTensors first. Found in aten.copy_.default

[anijain2305]

Related - #1447