fix chunk move device (#1158)

2 years ago · 4e67b2a890
parent 07f9c781f9
commit 4e67b2a890
1 changed files with 1 additions and 1 deletions
--- 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