I am trying to build a model that iterates two different layers, and below is my implementation.
class mymodel(nn.Module):
    def __init__(self, k_iters):
        super().__init__()
        self.k_iters = k_iters
        self.layer1 = ..
        self.layer2 =..

    def forward(self, x):
        x_k = x
        for k in range(self.k_iters):
            z_k = self.layer1(x_k)
            x_k = self.layer2(z_k)
        return x_k

Would there be any problem when I reuse the variables “z_k” and “x_k” for many times? I wonder if it interferes with the backward operation or tracking gradient…