Sooo as it happens from the text side.
Can you check if any output contains NaNs running the text model under with torch.no_grad() and text model in eval mode?
(This is textmode.eval() )
(Obviusly since gradients won’t be computed you will have to remove text model from the optimizer)
The idea in behind is to check if some batch normalization layer is exploiding while training.
EDIT: yeh the typical iisue is that once there is a NaN wherever in the forward, it will infect gradients and weights due to backprop.
Other options are, does the model contain dangerous layers like divisions logarithms etcetera…
It would be nice if you can detect which layer is generating the NaN inside your model.
(You can use forward hooks to get intermidiate outputs)