It might not be exactly a drop in replacement (and my implementation follows a GTC presentation of the CuDNN developers more than what Baidu did), but they implement the same loss.
Two caveats

You need PyTorch 1.1’s zero_infinity=True flag if you wish to zero losses and gradients of “impossible” samples (where the probability of producing the target using the input is 0 because the target is too long).
There still is a bug that you won’t get gradients for 0 length targets on the GPU.

Most people who do the first correctly seem to get their training to work.
Best regards
Thomas