test_new_function.py

import torch
from torch.autograd import Variable, Function

class MultiplyAdd(Function):
  @staticmethod
  def forward(ctx, input1, scalar, input2):
    ctx.scalar = scalar
    ctx.test = {'test': 3}
    return input1 + scalar * input2
    
  @staticmethod
  def backward(ctx, grad_output):
    print(ctx.test)
    return grad_output, None, ctx.scalar * grad_output


a = Variable(torch.randn(3, 5), requires_grad=True)
b = Variable(torch.randn(3, 5), requires_grad=True)
scalar = 2
ctx = {}

result = MultiplyAdd().apply(a, scalar, b)
g = Variable(torch.ones(1), requires_grad=True)
result.sum().backward(g)