Training LSTM neural network with missing values in target data - error optim.step()

I am would like to train the LSTM neural network with missing values in the target data and user-defined loss function. However, there is an error after optim.step(), some weights/biases are nan. Is there any hints on this? Thanks

import torch
import numpy as np
from torch import nn

# Define a simple lstm model
class myLSTM(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(myLSTM, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size)
    def forward(self, input):
        output, _ = self.lstm(input)
        return output

# Input and target
input = torch.randn(10, 5, requires_grad=True)
target = torch.randn(10, 5)

# There is one missing values in the target data
target[0,0] = np.nan

# Create model
lstmModel = myLSTM(5, 5) 

# Loss function, optimization
def loss_function(y_true, y_predict):
    return torch.nanmean((y_true-y_predict)**2)

optim = torch.optim.Adam(lstmModel.parameters(), lr=0.01)

# Training with only 1 epoch
output = lstmModel(input)
optim.zero_grad()
error = loss_function(target, output)
error.backward()
optim.step()

lstmModel.state_dict()

One solution can be to mask the nan elements, try the following:

def loss_function(y_true, y_predict):
    mask = ~torch.isnan(y_true)
    return torch.mean((y_true[mask] - y_predict[mask])**2)