lars
25718f175e
Removes unused imports and an ambiguous variable name, narrows Optional types before use so ty's flow analysis is satisfied, swaps sum() over polars expressions for pl.sum_horizontal to avoid the Literal[0] fallback type, and converts numpy bin edges to plain lists before passing to matplotlib's hist (whose stub only accepts Sequence[float]). Also applies ruff format across the repo, which had drifted out of sync with the formatter. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
20 lines
675 B
Python
20 lines
675 B
Python
import torch
|
|
import torch.version
|
|
|
|
print(f"PyTorch version: {torch.__version__}")
|
|
print(f"CUDA available: {torch.cuda.is_available()}")
|
|
|
|
if torch.cuda.is_available():
|
|
print(f"CUDA version: {torch.version.cuda}")
|
|
print(f"Device count: {torch.cuda.device_count()}")
|
|
print(f"Device name: {torch.cuda.get_device_name(0)}")
|
|
|
|
# Run a small tensor op on the GPU
|
|
a = torch.randn(1000, 1000, device="cuda")
|
|
b = torch.randn(1000, 1000, device="cuda")
|
|
c = a @ b
|
|
torch.cuda.synchronize()
|
|
print(f"Matrix multiply: OK (result shape {c.shape}, device {c.device})")
|
|
else:
|
|
print("No CUDA device found — check driver/CUDA installation.")
|