inputs = x if isinstance(x, list) else [x]
x = inputs[-1]
for i in range(self.num_stages):
for letter in self.decoder_layout:
if letter in ["b"]:
x = self.decoder_b[b_counter](x)
b_counter += 1
elif letter in ["u"]:
x = self.decoder_u[u_counter](x)
u_counter += 1
After Change
for letter, layer in zip(self.layout, self.values()):
if letter in ["b", "u"]:
x = layer(x)
elif letter in ["c"] and self.skip and (i < len(inputs) - 2):
x = layer([x, inputs[-i - 3]])
i += 1
return x