How to handle done agents with QMIX

How severe does this issue affect your experience of using Ray?

  • Low: It annoys or frustrates me for a moment.

Hello,
I was trying to use QMix in a custom environment. Here the agents may end (done==True) before the episode ends.
I saw that QMix allows us to use a mask, but it is used to mask the actions of a single agent.
I was thinking about adding an additional mask to mask out the Q-values of the done agents. Do you think this is a good way to do it or there are better options?