def METHOD_NAME(self, x: Tensor) -> List[Tensor]:
    out = list()
    if self.is_rewritten:
        xs = self.net(x)
        for i, x in enumerate(xs):
            if i in self.return_index:
                out.append(x)
    else:
        for i, cell in enumerate(self.net.cell_list):
            x = cell(x)
            if i in self.return_index:
                out.append(x)
    return out