From ee634beb2b2290cb614a855321af627f429d1ae3 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Tue, 30 Nov 2021 19:11:00 +0800
Subject: [PATCH] fix(pylite): fix device_id to dev_id

GitOrigin-RevId: 7989fe7dace86ffdb67e518cf92214342a53ec10
---
 lite/pylite/megenginelite/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lite/pylite/megenginelite/utils.py b/lite/pylite/megenginelite/utils.py
index 7b6110595..670aa910b 100644
--- a/lite/pylite/megenginelite/utils.py
+++ b/lite/pylite/megenginelite/utils.py
@@ -32,7 +32,7 @@ class TensorBatchCollector:
         self._mutex = threading.Lock()
         self.dev_type = device_type
         self.is_pinned_host = is_pinned_host
-        self.dev_id = 0
+        self.dev_id = device_id
         self.shape = shape
         self.dtype = LiteLayout(dtype=dtype).data_type
         self._free_list = list(range(self.shape[0]))
-- 
GitLab