li-plus's picture
Fix hidden_states dtype matching
2caddb7