Skip to content

Commit

Permalink
Merge pull request #234 from sensein/232-bug-_select_device_and_dtype…
Browse files Browse the repository at this point in the history
…-does-not-work-on-virtualized-hardware

Test Device Truly Available by Creating Trivial Object
  • Loading branch information
fabiocat93 authored Jan 22, 2025
2 parents d4da30a + 12904d3 commit 62117c1
Showing 1 changed file with 12 additions and 2 deletions.
14 changes: 12 additions & 2 deletions src/senselab/utils/data_structures/device.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,12 +44,22 @@ def _select_device_and_dtype(
if user_preference:
if not isinstance(user_preference, DeviceType):
raise ValueError(f"user_preference should be of type DeviceType, not {type(user_preference)}")

available_devices = [DeviceType.CPU]

if torch.cuda.is_available():
available_devices.append(DeviceType.CUDA)
try:
torch.empty(0, device=DeviceType.CUDA.value)
available_devices.append(DeviceType.CUDA)
except Exception as e:
print(f"CUDA is available but encountered an error: {e}")

if torch.backends.mps.is_available():
available_devices.append(DeviceType.MPS)
try:
torch.empty(0, device=DeviceType.MPS.value)
available_devices.append(DeviceType.MPS)
except Exception as e:
print(f"MPS is available but encountered an error: {e}")

# Check compatible and available
useable_devices = []
Expand Down

0 comments on commit 62117c1

Please sign in to comment.