From b570ad4b90fc9eb01d9b11f1c21314f3521c0bdf Mon Sep 17 00:00:00 2001 From: Yong He Date: Mon, 8 Jan 2024 09:58:03 -0800 Subject: Update a1-02-slangpy.md --- docs/user-guide/a1-02-slangpy.md | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) (limited to 'docs/user-guide') diff --git a/docs/user-guide/a1-02-slangpy.md b/docs/user-guide/a1-02-slangpy.md index b704e215a..bebdb06f4 100644 --- a/docs/user-guide/a1-02-slangpy.md +++ b/docs/user-guide/a1-02-slangpy.md @@ -201,7 +201,7 @@ class MySquareFunc(torch.autograd.Function): Now we can use the autograd function `MySquareFunc` in our python script: ```python -x = torch.tensor([[3.0, 4.0],[0.0, 1.0]], requires_grad=True, device='cuda') +x = torch.tensor((3.0, 4.0), requires_grad=True, device='cuda') print(f"X = {x}") y_pred = MySquareFunc.apply(x) loss = y_pred.sum() @@ -211,10 +211,9 @@ print(f"dX = {x.grad.cpu()}") Output: ``` -X = tensor([[3., 4.], - [0., 1.]], device='cuda:0', requires_grad=True) -dX = tensor([[6., 8.], - [0., 2.]]) +X = tensor([3., 4.], + device='cuda:0', requires_grad=True) +dX = tensor([6., 8.]) ``` And that's it! `slangpy.loadModule` uses JIT compilation to compile your Slang source into CUDA binary. -- cgit v1.2.3