def __call__(self, x, *args, **kwargs):
    features = self.tail(x, *args, **kwargs)
    print(features.shape)
    return (self.head(features, *args, **kwargs), features)
def __call__(self, x, *args, **kwargs):
    features = self.tail(x, *args, **kwargs)
    print(features.shape)
    return (self.head(features, *args, **kwargs), features)