def get_max_memory_allocated(device: str='cuda'):
    """
    Finds out the maximum memory allocated, then clears the max memory allocated.

    This currently only works for PyTorch models.

    TODO: Support TensorFlow and MXNet.

    :param device: (str) - name of device (Torch style) -> e.g. 'cuda'
    :return: float or None - if torch is in the environment, max memory allocated, else None
    """
    try:
        import torch
        max_mem = torch.cuda.max_memory_allocated(device=device)
        torch.cuda.reset_max_memory_allocated(device=device)
        return max_mem
    except ImportError:
        return None
