I don’t know how the model is updated exactly, but you could check all .grad attributes after the loss.backward() call and see, if there are nan values.
Generally, a nan loss could break your model as seen here:
torch.manual_seed(2809)

x = torch.randn(10, requires_grad=True)
optimizer = torch.optim.SGD([x], lr=1.)

y = torch.log(x)
y = y * y
loss = y.mean()
print(loss)
> tensor(nan, grad_fn=<MeanBackward0>)

loss.backward()
print(x.grad)
> tensor([    nan, -0.0501,     nan, -0.0420,     nan,     nan,     nan,     nan,
        -0.2245,     nan])

optimizer.step()
print(x)
> tensor([   nan, 0.8653,    nan, 0.8805,    nan,    nan,    nan,    nan, 0.7679,
           nan], requires_grad=True)

Are you sure the optimizer.step() is indeed performed or does your code have specific guards to avoid it?