Closure scoping for nested objax.Functions about objax HOT 2 CLOSED

models = []
for i in range(2):
    models.append(objax.nn.Linear(2, 3))

def loss(x, y, m):
    return ((m(x) - y) ** 2).mean()

new_funcs = []
for m in models:
    new_funcs.append(objax.Function(lambda x, y, model=m: loss(x, y, model), m.vars()))

models = []
for i in range(2):
    models.append(objax.nn.Linear(2, 3))

new_funcs = []
for m in models:
    def loss(x, y, model=m):
        return ((model(x) - y) ** 2).mean()
    new_funcs.append(objax.Function(loss, m.vars()))