Skip to content

Commit 3e6528c

Browse files
authored
Tsingmicro chip count (flagos-ai#3887)
* Add .dockerignore file * Support chip count for TsingMicro backend
1 parent 4683414 commit 3e6528c

2 files changed

Lines changed: 14 additions & 0 deletions

File tree

.dockerignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
.venv

tools/run_tests.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -232,6 +232,19 @@ def _probe_torch():
232232
pinfo(f"PyTorch device count ... {dev_count}")
233233
except Exception:
234234
ENV_INFO["torch"]["device_count"] = 0
235+
dev_count = 0
236+
237+
if dev_count == 0:
238+
try:
239+
# Is this a TsingMicro chip?
240+
import torch_txda
241+
242+
dev_count = torch_txda.device_count()
243+
244+
ENV_INFO["torch"]["device_count"] = dev_count
245+
pinfo(f"TorchTXDA device count ... {dev_count}")
246+
except Exception:
247+
pass
235248

236249

237250
def _probe_triton():

0 commit comments

Comments
 (0)