Is there an alterntive to tf.keras.layers.Masking in PyTorch? Or Is there any way to implement it from scratch?

I have a time series dataset with a lot of NAs that I need to use with LSTM network. Previously with TensorFlow, I used to initially replace NAs with -1(Which is not present in the data) and use tf.keras.layers.Masking(Documentation) within the model to stop learning when the model encounters -1 and resume when encountering something else. Since then, I have switched to PyTorch and need to use something similar again. So, my question is this: Is there an alternative to this in PyTorch? Or is there any way to do this while using PyTorch? See the code version of what I do below:

...
dataset.fillna(-1, inplace = True) # Replace na cells with -1 for masking in LSTM
...
...
model = tf.keras.models.Sequential()
model.add(tf.keras.layers.Masking(mask_value=-1,input_shape=(timesteps, features)))
model.add(tf.keras.layers.LSTM(32))
...
...
1 Like