def METHOD_NAME(
    diar_config_file2,
    batch_size,
    input_size,
    segment_size,
    normalize_segment_scale,
    num_spk,
):
    diarize_speech = DiarizeSpeech(
        train_config=diar_config_file2,
        segment_size=segment_size,
        normalize_segment_scale=normalize_segment_scale,
        num_spk=num_spk,
    )
    wav = torch.rand(batch_size, input_size)
    diarize_speech(wav, fs=8000)