def METHOD_NAME(*args, **kwargs):
    if args[3] in ["train", "test", "valid", "test_1", "test_2"]:
        data_path = args1.joinpath(f"rnn_input_{args[3]}.pkl")
        kwargs['data_path'] = data_path
    if args2:
        dataset = torch.load(data_path)
        return dataset
    return func(*args, **kwargs)