multinomial grad changed to zeros_like