.png)
3 months ago
3
from chisel import ChiselApp, GPUType
# here we're choosing 1xA100-80GB
# go to /configuration to learn about GPU types
app = ChiselApp("my-app", gpu=GPUType.A100_80GB_1)
@app.capture_trace(trace_name="gpu_task")
def gpu_function():
import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
# Your GPU code here
x = torch.randn(1000, 1000, device=device)
y = torch.randn(1000, 1000, device=device)
return torch.mm(x, y).cpu().numpy()
if __name__ == "__main__":
result = gpu_function()
print(f"Result shape: {result.shape}")
