From c0a5ed2490bd7d08222d13c136a4a9615c291ad9 Mon Sep 17 00:00:00 2001
From: "Lai, Yejing" <yejing.lai@intel.com>
Date: Mon, 24 Feb 2025 01:00:06 -0800
Subject: [PATCH] Fix meta load tensor imcompatible issue

Signed-off-by: Lai, Yejing <yejing.lai@intel.com>
---
 deepspeed/module_inject/layers.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/deepspeed/module_inject/layers.py b/deepspeed/module_inject/layers.py
index 2a24c2920466..5087d71a3d62 100644
--- a/deepspeed/module_inject/layers.py
+++ b/deepspeed/module_inject/layers.py
@@ -48,7 +48,8 @@ def move(tensor, device):
     # to save host resources when DP > 1。
 
     if tensor.is_meta:
-        return torch.empty_like(tensor, device=device)
+        # Keep tensor in meta device if tensor is meta.
+        return tensor
     else:
         # Using new tensors help in freeing memory (after split for example) was done before by calling clone().
         # Using copy=True instead of clone() will help in case of cpu --> cpu.