import torch
from torch import nn


def select_device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


device = select_device()
model = nn.Linear(3, 2).to(device)
inputs = torch.ones(4, 3, device=device)

with torch.no_grad():
    outputs = model(inputs)

print(f"torch_version={torch.__version__}")
print(f"cuda_available={torch.cuda.is_available()}")
print(f"mps_available={torch.backends.mps.is_available()}")
print(f"selected_device={device}")
print(f"model_device={next(model.parameters()).device}")
print(f"input_device={inputs.device}")
print(f"output_device={outputs.device}")