Error with LSTM Layer on MPS

mps
1

Hi,

I am trying to implement an LSTM network and set it so calculations are on the GPU with .to(device) but I keep getting the same error:

NotImplementedError: The operator ‘aten::mkldnn_rnn_layer’ is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on General MPS op coverage tracking issue · Issue #77764 · pytorch/pytorch · GitHub. As a temporary fix, you can set the environment variable PYTORCH_ENABLE_MPS_FALLBACK=1 to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.

I thought that the ability to use MPS with an AMD GPU was already implemented? Does anyone know the issue?

Here is the code that produced the error:

x = torch.ones(1,1,10)

device = torch.device('mps')
x.to(device)

lstm_layer = torch.nn.LSTM(
    input_size=10,
    hidden_size=10,
    num_layers=1,
    batch_first=True
)

lstm_layer.to(device)

out, (h, c) = lstm_layer(x)

Versions:

  • MacOS Ventura 13.4.1, Intel processor
  • AMD Radeon Pro 5300M
  • pytorch 2.0.1, installed using anaconda