README.md 314 Bytes
Newer Older
zhangshao's avatar
zhangshao committed
1
2
3
4
5
6
7
8
9
This CUDA extension implements optimized cross-entropy loss, adapted from Apex's
[Xentropy](https://github.com/NVIDIA/apex/tree/master/apex/contrib/xentropy).
We make it work for bfloat16 and support in-place backward to save memory.

It has only been tested on A100s.

```sh
cd csrc/xentropy && pip install .
```