From c35bc1f296701c5e742d25fa3f01e02404b961d0 Mon Sep 17 00:00:00 2001 From: Hardik Sharma Date: Wed, 18 Dec 2024 10:12:38 -0800 Subject: [PATCH] Support torch.Tensor in GraphBuilder. Summary: Add support for using torch.Tensor instead of torch.FakeTensor in GraphBuilder. This avoids errors related to torch.Tensor when building graphs. Differential Revision: D67350487 --- backends/cadence/aot/graph_builder.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/backends/cadence/aot/graph_builder.py b/backends/cadence/aot/graph_builder.py index 27604eac321..88ed2ac7699 100644 --- a/backends/cadence/aot/graph_builder.py +++ b/backends/cadence/aot/graph_builder.py @@ -38,7 +38,10 @@ def __init__(self) -> None: self.tracer: ExportPass.ExportTracer = self.ExportTracer( self, torch.fx.graph.CodeGen() ) - self.fake_tensor_mode = FakeTensorMode(allow_fallback_kernels=False) + self.fake_tensor_mode = FakeTensorMode( + allow_fallback_kernels=False, + allow_non_fake_inputs=True, + ) self.tracer.fake_tensor_mode = self.fake_tensor_mode # This will be called to create nodes in tracer.