[PyTorch] torch.exp 와 auto mixed precision

모델을 구성하는 레이어(혹은 연산) 중 torch.exp 를 포함하는 경우 auto mixed precision과 충돌이 발생할 수 있습니다.

 

이는 torch.exp의 output data type이 auto mixed precision과 호환되도록 자동으로 변환되지 않아 발생하는 문제입니다. (추후 업데이트에서 해결될 문제로 보이긴 합니다.)

 

이런 경우 torch.exp 의 output의 data type을 입력 텐서 x의 data type으로 재변환시키는 과정이 필요합니다.

output = torch.exp(x).type(x.dtype)