Edit: Disregard all this, we figured out a workaround on our end.

@driazati @zdevito As far as I can tell, you removed 'LSTM' from _rnn_impls https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/rnn.py#L15-L19 because that layer of indirection isn't necessary anymore for LSTM (now that LSTM has its own full-fledged forward and doesn't need to lean on the forward inherited from RNNBase).

Unfortunately for us, _rnn_impls dict is what Amp interposes on to insert its arg-casting wrapper (we can't patch _VF.* because that comes directly from C++). This means that Amp is broken with for LSTMs with current master. Is it ok if LSTM continues to use the layer of indirection, i.e., restore

_rnn_impls = {
    'LSTM': _VF.lstm,
    'GRU': _VF.gru,
    'RNN_TANH': _VF.rnn_tanh,
    'RNN_RELU': _VF.rnn_relu,
}

and change the _VF.lstm() calls to _rnn_impls['LSTM']()? I realize this is not necessary for you, and in support of our own use case, but it is helpful to us and doesn't seem to do any harm. If you don't object I can PR it.

@mcarilli The change was due to some technical limitations in the JIT since it uses the same module code, good you found a fix since we likely wouldn't be able to fix the limitation for a few weeks