Do not apply fp16 to certain key returned by Dataset

[MattYoon]

My Dataset returns additional doc_id key which is passed to the model to log loss on each document individually within a batch. The problem is that doc_id has some large numbers over > 1000. Due to numerical precision issue with fp16, some doc_id will have wrong values when I access them during validation_step. When running with full-precision, the doc_ids are all correct.

Is there anyway I can avoid applying fp16 to doc_id? This shouldn't be a problem with the model as doc_id is never used to update the model's parameter. With vanilla PyTorch I think this should be pretty straight forward, however with Lightning it is unclear exactly when the fp16 is getting applied.

Thank you in advance.

[akihironitta]

@MattYoon Is each doc_id converted to float16 at the very beginning of your training_step?

[MattYoon]

Sorry for the late reply! I tried converting the doc_id to other data types including float16, but the result was the same.

[MattYoon]

I just found out that this happened because I tried to stack float16 doc_id with float32 loss to a single tensor. The doc_id values were actually correct in validation_step, it was corrupted when the stack operation was applied.