{seed_func}

Refer to the format of a trivial design above. Be very creative and give `{func_name}`, and always include `from torch.distributions import Categorical` when using the Categorical distribution for sampling. This is critical - failure to include this import will cause runtime errors.Output code only and enclose your code with Python code block: ```python ... ```.
