def __call__(self, x):
outer_dense = nn.jit(Outer)(name='outer')
# we share stateful outer module as arg to two different, transformed modules:
inner = nn.jit(Inner)(outer_dense, name='inner1')
inner2 = nn.jit(Inner)(outer_dense, name='inner2')
res = inner(x) + inner2(x)
return res
def __call__(self, x):
cntr = nn.jit(Counter)(name='cntr')()
return x
