Function init_zs

train_acc.py:86–104

(args, device, in_channels, input_size)

def init_zs(args, device, in_channels, input_size):
    if is_video(args):
        _zs = torch.randn(
            args.data.sample_fid_bs,
            args.model.params.video_frames,
            in_channels,
            input_size,
            input_size,
            device=device,
        )
    else:
        _zs = torch.randn(
            args.data.sample_fid_bs,
            in_channels,
            input_size,
            input_size,
            device=device,
        )
    return _zs


#################################################################################