Upload fractus/nn/ce.py with huggingface_hub
Browse files- fractus/nn/ce.py +3 -1
fractus/nn/ce.py
CHANGED
|
@@ -56,7 +56,9 @@ def chunked_cross_entropy(
|
|
| 56 |
for i in range(0, n, ce_chunk):
|
| 57 |
hi = h[i : i + ce_chunk]
|
| 58 |
ti = targets[i : i + ce_chunk]
|
| 59 |
-
|
|
|
|
|
|
|
| 60 |
part = checkpoint(_chunk_ce, hi, weight, ti, use_reentrant=False)
|
| 61 |
else:
|
| 62 |
part = _chunk_ce(hi, weight, ti)
|
|
|
|
| 56 |
for i in range(0, n, ce_chunk):
|
| 57 |
hi = h[i : i + ce_chunk]
|
| 58 |
ti = targets[i : i + ce_chunk]
|
| 59 |
+
# Single-chunk + non-reentrant checkpoint collides with a second
|
| 60 |
+
# CE pass (scheduled sampling). Direct CE is cheap at N<=ce_chunk.
|
| 61 |
+
if torch.is_grad_enabled() and n > ce_chunk:
|
| 62 |
part = checkpoint(_chunk_ce, hi, weight, ti, use_reentrant=False)
|
| 63 |
else:
|
| 64 |
part = _chunk_ce(hi, weight, ti)
|