Hi, I have written pytorch custom function to achieve some kind of mathematic transform. But I come across a RuntimeError:

Inked1_LI.jpg764×126 47.8 KB

But my input is indeed a variable since I put the input as input = Variable(torch.randn(2,1,32,32)), I still can not figure out what’s wrong with my script.
And my script is the following:
class DHTtransform(Function):
	@staticmethod
	def forward(ctx, input, center):
		ctx.center = center
		ctx.save_for_backward(input)
		dht = DHT2d(input.data.numpy())
		if ctx.center:
			dht = np.fft.fftshift(dht, axes=(2,3,))
		dht = torch.from_numpy(dht)
		return Variable(dht)

	@staticmethod
	def backward(ctx, grad_output):
		input, = ctx.saved_variables
		grad_input = None
		if ctx.needs_input_grad[0]:
			dht = DHT2d(grad_output.data.numpy())
			if ctx.center:
				dht = np.fft.fftshift(dht, axes=(2,3,))
			dht = torch.from_numpy(dht)
			grad_input = Variable(dht)
		return grad_input

class DHTlayer(nn.Module):
	def __init__(self, center=False):
		super(DHTlayer, self).__init__()
		self.center = center

	def forward(self, input):
		return DHTtransform.apply(input, self.center)

Could someone help me out? Any advice would be appreciated.