Support many GPUs on PyTorch!

device_picker.py

# License:      Unlicense, attribution optional but appreciated
# Copyright:    Copyright (c) 2023 mrfakename. All rights reserved. Distributed under the Unlicense license.
# Author:       mrfakename
# Date created: Dec 23, 2023
# Published on: GitHub Gist

import torch

def get_device():
	if torch.cuda.is_available():
		return 'cuda'
	if torch.backends.mps.is_available():
		return 'mps'
	return 'cpu'

example_usage.py

# License:      Unlicense, attribution optional but appreciated
# Copyright:    Copyright (c) 2023 mrfakename. All rights reserved. Distributed under the Unlicense license.
# Author:       mrfakename
# Date created: Dec 23, 2023
# Published on: GitHub Gist

from device_picker import get_device

# ... #

model.to(get_device())