From 6b18293968a32cdcb2d9051f3efe421fa1335994 Mon Sep 17 00:00:00 2001
From: Wenhao Chen <cwher@outlook.com>
Date: Tue, 28 Nov 2023 16:06:08 +0800
Subject: [PATCH] to: add TODO mark

---
 colossalai/utils/memory.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/colossalai/utils/memory.py b/colossalai/utils/memory.py
index 955a588a1e28..efe4b4f28173 100644
--- a/colossalai/utils/memory.py
+++ b/colossalai/utils/memory.py
@@ -54,6 +54,7 @@ def colo_device_memory_capacity(device: torch.device) -> int:
     Returns:
         int: size in byte
     """
+    # TODO: add NPU support
     assert isinstance(device, torch.device)
     if device.type == "cpu":
         # In the context of 1-CPU-N-GPU, the memory capacity of the current process is 1/N overall CPU memory.