Hi,
Could you try adding torch.set_flush_denormal(True) at the top of your script to see if that fixes the issue?
torch.set_flush_denormal(True)