From 55c800e632f18f8904a806d461188639aff94c35 Mon Sep 17 00:00:00 2001 From: Philippe Tillet Date: Fri, 13 Mar 2020 11:42:43 -0400 Subject: [PATCH] [DOCS] Fixed typo: triton.function -> torch.autograd.Function --- docs/tutorials/custom-operation.rst | 2 +- docs/tutorials/putting-it-all-together.rst | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/tutorials/custom-operation.rst b/docs/tutorials/custom-operation.rst index 30619bef9..e6abffb12 100644 --- a/docs/tutorials/custom-operation.rst +++ b/docs/tutorials/custom-operation.rst @@ -39,7 +39,7 @@ As you will see, a wrapper for the above Triton function can be created in just import torch import triton - class _add(triton.function): + class _add(torch.autograd.Function): # source-code for Triton compute kernel src = """ __global__ void add(float* z, float* x, float* y, int N){ diff --git a/docs/tutorials/putting-it-all-together.rst b/docs/tutorials/putting-it-all-together.rst index 3303750a6..4f760c83c 100644 --- a/docs/tutorials/putting-it-all-together.rst +++ b/docs/tutorials/putting-it-all-together.rst @@ -15,7 +15,7 @@ The PyTriton API provides a :code:`triton.function` class which automatically ha import triton # Entry point - class _dot(triton.function): + class _dot(torch.autograd.Function): @staticmethod # Forward Pass @@ -170,7 +170,7 @@ Creating custom operations for Triton and PyTorch is very similar; programmers h return da, db, None, None, None, None, None, None, None -A callable operation can be created using the :code:`apply` method of our :code:`triton.function` class. +A callable operation can be created using the :code:`apply` method of the :code:`torch.autograd.Function` class. .. code:: python