def __init__(self, dim: int, unflattened_size: Size) -> None:
        super(Unflatten, self).__init__()
        self.dim = dim
        self.unflattened_size = unflattened_size

    def forward(self, input: Tensor) -> Tensor:
        return input.unflatten(
            self.dim,
            (('C', self.unflattened_size[0]), ('H', self.unflattened_size[1]), ('W', self.unflattened_size[2]))
        )

class UnFlatten(nn.Module)
    def __init__(self, dim: Union[int,str], unflattened_size: Union[tuple,Size]) -> None:
        super(Unflatten, self).__init__()
        if isinstance(dim, int):
            self.named = False
            # Make sure the unflattened_size is a tuple of ints
        else:
            self.named = True
            # Make sure the unflattened_size is a tuple of tuple

        self.dim = dim
        self.unflattened_size = unflattened_size

    def forward(self, input: Tensor) -> Tensor:
        if self.named:
            return input.unflatten(self.dim, self.unflattened_size)
        else:
            dim = self.dim
            if self.dim < 0:
                dim += input.ndim()
            inp_size = list(input.size())
            new_size = inp_size[:dim] + self.unflattened_size + inp_size[dim+1:]
            return input.view(new_size)