import torch
import pickle


def load_model(model_pkl_path, device=None):
    """Load StyleGAN2 model from pickle checkpoint file."""

    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    with open(model_pkl_path, 'rb') as f:
        G = pickle.load(f)['G_ema'].to(device)  # torch.nn.Module
    return G