Thanks for the answer! May I know if the masked entry will affect the gradient? For example, if I have a model whose intermediate layer gives an output of:
out = [inf, -3.4, inf, -5.5]

where the inf entries should be masked out, can I do it by:
out[out == float('inf')] = 0

Thanks!