in jcm/models/layers.py [0:0]
def __call__(self, xs):
rcu_block = functools.partial(RCUBlock, n_blocks=2, n_stages=2, act=self.act)
rcu_block_output = functools.partial(
RCUBlock,
features=self.features,
n_blocks=3 if self.end else 1,
n_stages=2,
act=self.act,
)
hs = []
for i in range(len(xs)):
h = rcu_block(features=xs[i].shape[-1])(xs[i])
hs.append(h)
if not self.start:
msf = functools.partial(
MSFBlock, features=self.features, interpolation=self.interpolation
)
h = msf(shape=self.output_shape)(hs)
else:
h = hs[0]
crp = functools.partial(
CRPBlock, features=self.features, n_stages=2, act=self.act
)
h = crp()(h)
h = rcu_block_output()(h)
return h