Error about hooks in implementing customized modules
Hi, I just implemented an customized modules as follows: <SCODE>def forward(self, x, para):
	''' 
	Implement eq 26 of ref paper for each batch.
	Input:
		para: dim(seq_len, batch, 121)
		x:    dim(seq_len, batch, 3)
	'''
	if x.size()[0] == para.size()[0]:
		self.seq_length = x.size()[0] - 1
		total_loss = 0
		for i in range(self.seq_length):
			# prepare parameters
			self.__get_para(i, x, para)
			normalpdf = self.__para2normal(self.x1, self.x2, self.m1, self.m2, self.s1, self.s2, self.rho) #dim (n_batch, 20)
			single_loss = self.__singleLoss(normalpdf)
			total_loss += single_loss
		return total_loss
	else:
		raise Exception("x and para don't match")


def __get_para(self, i, x, para):
	'''
	Slice and process parameters to the right form.
	Implementing eq 18-23 of ref paper.
	'''
	self.batch = x.size()[1]
	self.e = torch.sigmoid(-para[i,:,0])  # eq 18
	self.parameters = para

	# slice remaining parameters and training inputs
	self.pi, self.m1, self.m2, self.s1, self.s2, self.rho = torch.split(self.parameters[i,:,1:], 20, dim = 1) # dim(batch, 20)
	self.x1 = x[i+1,:,0].resize(self.batch, 1)  # dim(batch, 1)
	self.x2 = x[i+1,:,1].resize(self.batch, 1)
	self.et = x[i+1,:,2].resize(self.batch, 1)
	
	## process parameters
	# pi
	max_pi = torch.max(self.pi, dim = 1)[0]
	max_pi = max_pi.expand_as(self.pi)
	diff = self.pi - max_pi
	red_sum = torch.sum(diff, dim = 1).expand_as(self.pi)
	self.pi = diff.div(red_sum)

	# sd
	self.s1 = self.s1.exp()
	self.s2 = self.s2.exp()

	# rho
	self.rho = self.rho.tanh()


	# reshape ground truth x1, x2 to match m1, m2 because broadcasting is currently not supported by pytorch
	self.x1 = self.x1.expand_as(self.m1)
	self.x2 = self.x2.expand_as(self.m2)


def __para2normal(self, x1, x2, m1, m2, s1, s2, rho):
	'''
	Implement eq 24, 25 of ref paper.
	All input with dim(1, batch, 20)
	'''
	norm1 = x1.sub(m1)
	norm2 = x2.sub(m2)
	s1s2 = torch.mul(s1, s2)
	z = torch.pow(torch.div(norm1, s1), 2) + torch.pow(torch.div(norm2, s2), 2) - \
		2*torch.div(torch.mul(rho, torch.mul(norm1, norm2)), s1s2)
	negRho = 1 - torch.pow(rho, 2)
	expPart = torch.exp(torch.div(-z, torch.mul(negRho, 2)))
	coef = 2*np.pi*torch.mul(s1s2, torch.sqrt(negRho))
	result = torch.div(expPart, coef)
	return result

	
def __singleLoss(self, normalpdf):
	'''
	Calculate loss for single time stamp. eq 26
	Input: normalpdf (1,n_batch, 20).
	'''
	epsilon = 1e-20  # floor of loss from mixture density component since initial loss could be zero
	mix_den_loss = torch.mul(self.pi, normalpdf)
	red_sum_loss = torch.sum(torch.log(mix_den_loss))  # sum for all batch
	end_loss = torch.sum(torch.log(torch.mul(self.e, self.et) + torch.mul(1-self.e, 1 - self.et)))
	total_loss = -red_sum_loss - end_loss
	
	return total_loss/self.batch
<ECODE>
Shihan_su
Hi, <SCODE>class Loss(torch.nn.Module):
    def __init__(self):
        super(torch.nn.Module, self).__init__()

        ... # rest of the __init__ here
<ECODE>
elanmart