From 4e67b2a890ca6aeb2cdaa9c455af793070641519 Mon Sep 17 00:00:00 2001
From: ver217 <lhx0217@gmail.com>
Date: Wed, 22 Jun 2022 18:07:10 +0800
Subject: [PATCH] fix chunk move device (#1158)

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

diff --git a/colossalai/tensor/chunk.py b/colossalai/tensor/chunk.py
index b3cb07328..4cacb0c7b 100644
--- a/colossalai/tensor/chunk.py
+++ b/colossalai/tensor/chunk.py
@@ -455,7 +455,7 @@ class ChunkManager:
             chunk (Chunk): the chunk to move to target device
             device (torch.device): target device
         """
-        if chunk.data.device == device:
+        if chunk.device_type == device.type:
             return
         if chunk.can_move_device and not chunk.is_empty:
             self.total_mem[chunk.device_type] -= chunk.mem