import torch
from torch import nn


print(f"torch={torch.__version__}")
print(f"rocm_build={torch.version.hip}")
print(f"cuda_runtime={torch.version.cuda}")
print(f"cuda_available={torch.cuda.is_available()}")

if not torch.version.hip:
    raise SystemExit("This PyTorch install is not a ROCm build")
if not torch.cuda.is_available():
    raise SystemExit("ROCm GPU is not available to PyTorch")

device = torch.device("cuda")
x = torch.tensor([1.0, 2.0, 3.0], device=device)
y = x * 2
model = nn.Linear(3, 1).to(device)

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

print(f"device_count={torch.cuda.device_count()}")
print(f"device_name={torch.cuda.get_device_name(0)}")
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("rocm_smoke=True")