import platform

import torch
from torch import nn


print(f"python={platform.python_version()}")
print(f"torch={torch.__version__}")
print(f"machine={platform.machine()}")
print(f"mps_built={torch.backends.mps.is_built()}")
print(f"mps_available={torch.backends.mps.is_available()}")

if not torch.backends.mps.is_available():
    if not torch.backends.mps.is_built():
        raise SystemExit("MPS support is not built into this PyTorch install")
    raise SystemExit("MPS is built, but this macOS host has no available MPS device")

device = torch.device("mps")
x = torch.ones(3, device=device)
y = x * 2
model = nn.Linear(3, 1).to(device)

with torch.no_grad():
    output = model(y)

print(f"tensor={y}")
print(f"tensor_device={y.device}")
print(f"model_device={next(model.parameters()).device}")
print(f"output_device={output.device}")
print("mps_smoke=True")