From 5d3f7bc4b6c2505d6cc33e0e5116669c914e1e1c Mon Sep 17 00:00:00 2001
From: dainiu <19307110036@fudan.edu.cn>
Date: Fri, 20 Jan 2023 16:55:58 +0800
Subject: [PATCH] [hotfix] meta tensor default device.

---
 colossalai/fx/profiler/tensor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/colossalai/fx/profiler/tensor.py b/colossalai/fx/profiler/tensor.py
index 7606f17cf9d5..2ee5e5c47750 100644
--- a/colossalai/fx/profiler/tensor.py
+++ b/colossalai/fx/profiler/tensor.py
@@ -43,7 +43,7 @@ def __new__(cls, elem, fake_device=None):
             storage_offset=elem.storage_offset(),
             dtype=elem.dtype,
             layout=elem.layout,
-            device=fake_device if fake_device is not None else torch.device('cpu'),
+            device=fake_device or (elem.device if elem.device.type != 'meta' else torch.device('cpu')),
             requires_grad=elem.requires_grad)    # deceive the frontend for aten selections
         r._tensor = elem
         # ...the real tensor is held as an element on the tensor.