From fa36a84bf52056ff9d3edc11e8afb860057a3c87 Mon Sep 17 00:00:00 2001
From: Max Ren <maxren@meta.com>
Date: Tue, 16 Apr 2024 17:12:56 -0700
Subject: [PATCH] change call_delegate_autograd (#3073)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Summary:

Some changes angela told me to make 😂

Differential Revision: D56222503
---
 exir/delegate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/exir/delegate.py b/exir/delegate.py
index 959bd4bb17c..076e08daf37 100644
--- a/exir/delegate.py
+++ b/exir/delegate.py
@@ -102,7 +102,7 @@ def fake_requires_grad(var):
                         var.requires_grad = True
                 return var
 
-            return pytree.tree_map(fake_requires_grad, res)
+            return pytree.tree_map_only(torch.Tensor, fake_requires_grad, res)
 
         return res