def METHOD_NAME(self, x):
    x = self.features(x)
    if isinstance(x, tuple):
        x = x[0]  # it might be a namedtuple such as torchvision.model.InceptionOutputs
    elif torch.jit.isinstance(x, Dict[str, torch.Tensor]):
        x = x[self.node_name]  # torchvision create_feature_extractor
    if self.pool is not None:
        x = self.pool(x)
    if not self.use_conv:
        x = torch.flatten(x, 1)
    else:  # user specified `use_conv` but the pooling layer removed the spatial dims
        while len(x.shape) < self.dim + 2:
            x = x[..., None]
    x = self.fc(x)
    return x