pytorch: self defined layers

self_defined_layers.py

import torch


class MyFunction(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        output = torch.sign(input)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        # saved tensors - tuple of tensors, so we need get first
        input, = ctx.saved_variables
        grad_output[input.ge(1)] = 0
        grad_output[input.le(-1)] = 0
        return grad_output


# usage
x = torch.randn(10, 20)
y = MyFunction.apply(x)
# or
my_func = MyFunction.apply
y = my_func(x)


# and if we want to use inside nn.Module
class MyFunctionModule(torch.nn.Module):
    def forward(self, x):
        return MyFunction.apply(x)