Questions about calculation of t_k in IEM module.

[onpix]

As your paper writes:

But your code in model.py is:

def forward(self, input_y, input_u, k):
  if k == 0:
      t_hat = self.max_operation(input_y)
  else:
      t_hat = self.max_operation(input_u) - 0.5 * (input_u - input_y)
  t = t_hat
  t = self.cell(t)
  t = self.activate(t)
  t = torch.clamp(t, 0.001, 1.0)
  u = torch.clamp(input_y / t, 0.0, 1.0)

  return u, t

where t = self.cell(t_hat) instead of t = t_hat - self.cell(t_hat) as your paper shows.

Which one is correct? Looking forward to your reply.