silvester:

Does it also work when you wrap the model insider nn.sequential?


Yea, I think your error comes from some other place…
class M(nn.Module):
	def __init__(self):
		super(M, self).__init__()
		self.s = nn.Sequential(nn.ELU(), weightedReLU())
		
	
	def forward(self,x):
		return self.s(x)
	
m = M()
y=m.forward(torch.ones(5).requires_grad_())
y.sum().backward()

(your code untouched, except I removed cuda())



 silvester:

what’s the difference between defining a class as autograd function of as nn.Module?


Function is not related to nn.Module, you don’t even create instances of it explicitly, it is just a way to provide backward(). Module base class provides has all usual facilities for submodules, parameters, module tree enumeration etc.