Thank you Junpeng.
By “array_slice”, do you mean a single sample? For my model a single sample of accept is a six-element array, i.e. of shape (6,). logp.array_to_dict() seems to expect an array of shape (402,).
Am I missing something?