Skip to content

Commit

Permalink
Merge branch 'develop' into shrivaths/fix-maximum-tracking
Browse files Browse the repository at this point in the history
  • Loading branch information
talmo authored Sep 8, 2023
2 parents e1e7285 + 64655d6 commit ee26e2e
Showing 3 changed files with 10 additions and 0 deletions.
3 changes: 3 additions & 0 deletions sleap/nn/__init__.py
Original file line number Diff line number Diff line change
@@ -14,3 +14,6 @@
import sleap.nn.tracking
import sleap.nn.viz
import sleap.nn.identity
import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
1 change: 1 addition & 0 deletions sleap/nn/system.py
Original file line number Diff line number Diff line change
@@ -195,6 +195,7 @@ def get_gpu_memory() -> List[int]:
A list of the available memory on each GPU in MiB.
"""

if shutil.which("nvidia-smi") is None:
return []

6 changes: 6 additions & 0 deletions tests/nn/test_system.py
Original file line number Diff line number Diff line change
@@ -87,3 +87,9 @@ def test_gpu_order_and_length():

# Assert that the order and length of GPU indices match
assert sleap_indices == nvidia_indices


def test_gpu_device_order():
"""Indirectly tests GPU device order by ensuring environment variable is set."""

assert os.environ["CUDA_DEVICE_ORDER"] == "PCI_BUS_ID"

0 comments on commit ee26e2e

Please sign in to comment.