import torch

# Check if CUDA is available
[device](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "torch.device") = [torch.device](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "torch.device")('cpu')
if [torch.cuda.is_available](https://docs.pytorch.org/docs/stable/generated/torch.cuda.is_available.html#torch.cuda.is_available "torch.cuda.is_available")():
    [device](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "torch.device") = [torch.device](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "torch.device")('cuda')

[torch.set_default_device](https://docs.pytorch.org/docs/stable/generated/torch.set_default_device.html#torch.set_default_device "torch.set_default_device")([device](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "torch.device"))
print(f"Using device = {[torch.get_default_device](https://docs.pytorch.org/docs/stable/generated/torch.get_default_device.html#torch.get_default_device "torch.get_default_device")()}")