From 48753dffdf92ade402b33caa4a1c3530c2828503 Mon Sep 17 00:00:00 2001
From: hengyujiang <18550376097@163.cm>
Date: Sat, 16 Nov 2024 21:00:28 +0800
Subject: [PATCH] Fix: CustomTracer originally failed to trace Tensor object

---
 mmrazor/models/task_modules/tracer/fx/custom_tracer.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/mmrazor/models/task_modules/tracer/fx/custom_tracer.py b/mmrazor/models/task_modules/tracer/fx/custom_tracer.py
index 68d5f0809..41b654bc9 100644
--- a/mmrazor/models/task_modules/tracer/fx/custom_tracer.py
+++ b/mmrazor/models/task_modules/tracer/fx/custom_tracer.py
@@ -142,6 +142,10 @@ def _get_attrs(target, attrs):
             if isinstance(attr, nn.Module):
                 module_dict[node.target] = nn.Module()
                 special_nodes.append(node)
+            #; the original design fails to 
+            #; trace any Tensor object
+            elif isinstance(attr, torch.Tensor):
+                module_dict[node.target] = attr
         elif node.op == 'call_method':
             for special_node in special_nodes:
                 if special_node in node.args or \