# Ignore masked tokens (Flux.jl)

**URL:** <https://discourse.julialang.org/t/ignore-masked-tokens-flux-jl/75591>\
**Category:** Machine Learning\
**Created:** [February 1, 2022, 5:50pm UTC](https://discourse.julialang.org/t/ignore-masked-tokens-flux-jl/75591 "2022-02-01T17:50:59Z")\
**Posts on this page:** 7\
**Page:** 1

<div class="post-metadata">

**Author:** ![spragud2](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/spragud2/32/29409_2.png) [@spragud2](https://discourse.julialang.org/u/spragud2)\
**Post date:** [February 1, 2022, 5:50pm UTC](https://discourse.julialang.org/t/ignore-masked-tokens-flux-jl/75591/1 "2022-02-01T17:50:59Z")

</div>

Hi,

I’m working on implementing a seq2seq model in Flux. Is there a standard method, or if not a suitable approach, for ignoring masked tokens when computing the gradient? E.g. sentences of differing lengths, following the EOS token

Cheers

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [February 2, 2022, 4:27pm UTC](https://discourse.julialang.org/t/ignore-masked-tokens-flux-jl/75591/2 "2022-02-02T16:27:16Z")

</div>

I’m not aware of any existing functionality. If you have some insight into how other frameworks handle this, we could workshop something.

---

<div class="post-metadata">

**Author:** ![spragud2](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/spragud2/32/29409_2.png) [@spragud2](https://discourse.julialang.org/u/spragud2)\
**Post date:** [February 24, 2022, 7:45pm UTC](https://discourse.julialang.org/t/ignore-masked-tokens-flux-jl/75591/3 "2022-02-24T19:45:09Z")

</div>

It appears that PyTorch has an ignore\_index argument for its loss functions, which allows you to specify indices to ignore when computing the loss and will therefore affect gradient computations in a way I desire.

I guess it would be relatively straightforward to write my own loss function in Julia/Flux that does the same.

---

<div class="post-metadata">

**Author:** ![jling](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jling/32/212909_2.png) [@jling](https://discourse.julialang.org/u/jling)\
**Post date:** [February 24, 2022, 10:35pm UTC](https://discourse.julialang.org/t/ignore-masked-tokens-flux-jl/75591/4 "2022-02-24T22:35:07Z")

</div>

yeah, one of the “problems” of Julia is that it’s often too straightforward to write your own little something once you’re semi domain expert.

This means absolute beginners find Julia ecosystem harder to use due to less canned (although trivial) functionalities.

---

<div class="post-metadata">

**Author:** ![spragud2](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/spragud2/32/29409_2.png) [@spragud2](https://discourse.julialang.org/u/spragud2)\
**Post date:** [February 24, 2022, 11:34pm UTC](https://discourse.julialang.org/t/ignore-masked-tokens-flux-jl/75591/5 "2022-02-24T23:34:35Z")

</div>

Honestly this is so true its painful. TF/pytorch are functionally languages unto themselves, and I often feel frozen out of writing custom algo’s because I simply don’t know TF/pytorch well enough. Always left asking, am I writing this correctly?

With Julia/Flux, the code is the math, and everything just works. Love it!

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [March 4, 2022, 4:51pm UTC](https://discourse.julialang.org/t/ignore-masked-tokens-flux-jl/75591/6 "2022-03-04T16:51:23Z")

</div>

If you do come up with something, feel free to open a PR and we’ll try to find a good home for it 🙂

---

<div class="post-metadata">

**Author:** ![cirobr](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cirobr/32/219994_2.png) [@cirobr](https://discourse.julialang.org/u/cirobr)\
**Post date:** [April 24, 2026, 7:56pm UTC](https://discourse.julialang.org/t/ignore-masked-tokens-flux-jl/75591/7 "2026-04-24T19:56:14Z")

</div>

Cheers,

I am stuck with a somehow similar problem, hope someone can advise.

Have tried to mimic Torchvision’s crossentropy, which accepts a mask of integers as input. The function also has an ignore\_class argument which is useful to exclude classes with small representation that just add noise to the process (typically class 255). Long story short, I came up with an algorithm that masks out such pixels from the loss calculation.

Problem occurs when the training loop is precompiled. The loss function uses the Flux.onehotbatch() instruction that breaks during compilation. However, inference compiles and executes flawlessly.

My good friend Grok AI tells me that this is a known issue, and suggests switching the AD engine to Enzyme. With the new engine, compilation during training breaks for a different reason: excess of allocations.

Can anyone please confirm the issue with onehotbatch() and Zygote, and perhaps suggest a bypass?

Thanks in advance.