| from torch.autograd.function import Function | |
| class Softsign(Function): | |
| def forward(self, input): | |
| self.buffer = input.clone().abs_().add_(1) | |
| self.buffer_squared = False | |
| output = input.clone().div_(self.buffer) | |
| return output | |
| def backward(self, grad_output): | |
| if not self.buffer_squared: | |
| self.buffer.mul_(self.buffer) | |
| self.buffer_squared = True | |
| grad_input = grad_output.clone().div_(self.buffer) | |
| return grad_input | |