2022-11-14 13:52:00 +08:00
|
|
|
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.
|
2022-11-15 23:10:25 +08:00
|
|
|
|
|
|
|
|
It has only been tested on A100s.
|
|
|
|
|
|
2022-11-14 13:52:00 +08:00
|
|
|
```sh
|
|
|
|
|
cd csrc/xentropy && pip install .
|
|
|
|
|
```
|
