Ok, so I think that I found the cause of the bug here, where the int64 dtype (which is type_num 7) is forced by the step method instead of looking at the variables’ particular dtype (int32 by default). I’ll work on a patch soon but for now, setting on's dtype to int64 should do the trick