pytorch-info

import torch

#get info on torch + platform configuration
print(f"torch version: {torch.__version__}")
print(f"torch gpu available: {torch.cuda.is_available()}")
print(f"torch gpu count: {torch.cuda.device_count()}")

if torch.cuda.is_available():
  print(f"torch current: {torch.cuda.current_device()}")
else:
  print("torch current: cpu")