From 8c6386c364c159a3df74f982f3cff3a244815a31 Mon Sep 17 00:00:00 2001
From: James Chua <chuajamessh@gmail.com>
Date: Wed, 3 May 2023 23:27:02 +0800
Subject: [PATCH] add comment on future improvement

---
 elk/utils/multi_gpu.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/elk/utils/multi_gpu.py b/elk/utils/multi_gpu.py
index ab929484..f7cf11a3 100644
--- a/elk/utils/multi_gpu.py
+++ b/elk/utils/multi_gpu.py
@@ -42,6 +42,10 @@ def instantiate_model_with_devices(
     else:
         torch_dtype = "auto"
 
+    # TODO: Maybe we should ensure the device map is the same
+    # for all the extract processes? This is because the device map
+    # can affect performance highly and its annoying if one process
+    # is using a different device map than the others.
     device_map = (
         {"": first_device}
         if device_config.is_single_gpu