提交 313b519d authored 作者: Arnaud Bergeron's avatar Arnaud Bergeron

GpuDnnSoftmax works with float16.

上级 f3eb824d
...@@ -1353,6 +1353,7 @@ class GpuDnnSoftmax(GpuDnnSoftmaxBase): ...@@ -1353,6 +1353,7 @@ class GpuDnnSoftmax(GpuDnnSoftmaxBase):
or per spatial location '01' per image across 'c'. or per spatial location '01' per image across 'c'.
""" """
_f16_ok = True
direction = "forward" direction = "forward"
file = "dnn_softmax.c" file = "dnn_softmax.c"
c_func = "APPLY_SPECIFIC(softmax)" c_func = "APPLY_SPECIFIC(softmax)"
...@@ -1389,6 +1390,7 @@ class GpuDnnSoftmaxGrad(GpuDnnSoftmaxBase): ...@@ -1389,6 +1390,7 @@ class GpuDnnSoftmaxGrad(GpuDnnSoftmaxBase):
image across 'c'. image across 'c'.
""" """
_f16_ok = True
direction = 'backward' direction = 'backward'
file = "dnn_softmax_grad.c" file = "dnn_softmax_grad.c"
c_func = "APPLY_SPECIFIC(softmax_grad)" c_func = "APPLY_SPECIFIC(softmax_grad)"
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论