"softmax_*_kernel_impl" not implemented for 'Long' というエラーの解決法

エラー内容

Softmax 関数を用ようとした時，以下のようなエラーが発生しました．

RuntimeError: "softmax_lastdim_kernel_impl" not implemented for 'Long'

RuntimeError: "softmax_kernel_impl" not implemented for 'Long' 

問題が発生しなかった場合

softmax_01.py

m = nn.Softmax(dim=0)
input = torch.randn(2, 3, 4)
output = m(input)
# → Softmax によって axis=0 の方向で処理された出力

問題が発生した場合

softmax_02.py

m = nn.Softmax(dim=0)
input = torch.tensor([1, 2, 3, 4, 5])
output = m(input)
# → 表記のエラー

解決方法

エラー内容を読むとどうやら Softmax は Long タイプについて対応していなさそうなので，データの種類を変更してみたところ，解決しました．

softmax_03.py

m = nn.Softmax(dim=0)
input = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float64)
out = m(input)
# → Softmax によって axis=0 の方向で処理された出力