pytorch xla decode indices

[ma2za]

Hi, I am trying to use CRF with a model running on tpu, but in the _viterbi_decode function at the line 300

next_score, indices = next_score.max(dim=1)

I get these indices and next score is correct.

tensor([[-9223372036854775808, -9223372036854775808, -9223372036854775808,
          ..., -9223372036854775808, -9223372036854775808,
         -9223372036854775808],
        [-9223372036854775808,                   24, -9223372036854775808,
          ..., -9223372036854775808, -9223372036854775808,
         -9223372036854775808],
        [-9223372036854775808,                   20, -9223372036854775808,
          ..., -9223372036854775808,                   23,
                           23],
        ...,
        [-9223372036854775808,                   23, -9223372036854775808,
          ..., -9223372036854775808, -9223372036854775808,
         -9223372036854775808],
        [-9223372036854775808, -9223372036854775808, -9223372036854775808,
          ..., -9223372036854775808, -9223372036854775808,
         -9223372036854775808],
        [                  23, -9223372036854775808, -9223372036854775808,
          ..., -9223372036854775808, -9223372036854775808,
                           23]]) 

If I add the following command the line before I get the correct indices:

_ = next_score.cpu()

So I don't use the the cpu tensor, but it is enough to correct the behavior.

 tensor([[ 2,  2,  2,  ...,  2,  2,  2],
        [24, 24, 24,  ..., 24, 24, 24],
        [23, 20, 24,  ..., 23, 23, 23],
        ...,
        [23, 23, 23,  ..., 23, 23, 23],
        [ 8,  8,  8,  ...,  8,  8,  8],
        [23, 31, 31,  ..., 23,  2, 23]])

Have you faced this problem before or do you have any idea what is the cause?

[kmkurn]

Hmm that's weird. I've never had this issue before, but I've never tried running this on tpu either. I have no access to a tpu so I'm afraid I can't be much help here, sorry!