def METHOD_NAME(self, x, bases):
    """ Multiple updates of `bases` and `coeff` to better match `x`. """
    # (B * S, D, N)^T @ (B * S, D, R) -> (B * S, N, R)
    coeff = torch.bmm(x.transpose(1, 2), bases)
    coeff = F.softmax(self.inv_t * coeff, dim=-1)
    steps = self.n_train_steps if self.training else self.n_eval_steps
    for _ in range(steps):
        bases, coeff = self.local_step(x, bases, coeff)
    return bases, coeff