import torch


def squeeze_expand_dim(tensor, axis):
    """ helper to squeeze a multi-dim tensor and then
        unsqueeze the axis dimension if dims < 4"""
    tensor = torch.squeeze(tensor)
    if len(list(tensor.size())) < 4:
        return tensor.unsqueeze(axis)
    else:
        return tensor