simple script to figure out what pytorch device you are using

check-devices.py

import torch

if torch.cuda.is_available():
  print ("cuda available")
  print(f'available devices: {torch.cuda.device_count()}')
elif torch.backends.mps.is_available():
  print ("mps available")
  device = torch.device("mps")
  x = torch.ones(1, device=device)
  print (x)
else:
  print ("cpu availabile")

pyproject.toml

[tool.poetry]
name = "check-devices"
version = "0.1.0"
description = "simple script to figure out what pytorch device you are using"
authors = ["Jason Ross <[email protected]>"]
readme = "README.md"
package-mode = false

[tool.poetry.dependencies]
python = "^3.10"
torch = "^2.4.0"
numpy = "^2.1.0"


[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"