Browse Source

fix(pylite): fix device_id to dev_id

GitOrigin-RevId: 7989fe7dac
tags/v1.7.0.m1
Megvii Engine Team 3 years ago
parent
commit
ee634beb2b
1 changed files with 1 additions and 1 deletions
  1. +1
    -1
      lite/pylite/megenginelite/utils.py

+ 1
- 1
lite/pylite/megenginelite/utils.py View File

@@ -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]))


Loading…
Cancel
Save