Assert failing when using amp with rnns

My code: https://github.com/Rhuax/tesi/blob/master/train.py
The error:

Selected optimization level O2:  FP16 training with FP32 batchnorm and FP32 master weights.

Defaults for this optimization level are:
enabled                : True
opt_level              : O2
cast_model_type        : torch.float16
patch_torch_functions  : False
keep_batchnorm_fp32    : True
master_weights         : True
loss_scale             : dynamic
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                : True
opt_level              : O2
cast_model_type        : torch.float16
patch_torch_functions  : False
keep_batchnorm_fp32    : True
master_weights         : True
loss_scale             : dynamic
Warning:  multi_tensor_applier fused unscale kernel is unavailable, possibly because apex was installed without --cuda_ext --cpp_ext. Using Python fallback.  Original ImportError was: ModuleNotFoundError("No module named 'amp_C'",)
EPOCH:0
  0% 0/13604 [00:00<?, ?it/s]loading training set
loaded
Warning:  An input tensor was not cuda. 
Warning:  An input tensor was not cuda. 
Warning:  An input tensor was not cuda. 
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/rnn.py:179: RuntimeWarning: RNN module weights are not part of single contiguous chunk of memory. This means they need to be compacted at every call, possibly greatly increasing memory usage. To compact weights again call flatten_parameters().
  self.dropout, self.training, self.bidirectional, self.batch_first)


Traceback (most recent call last):
  File "train.py", line 424, in <module>
    train()
  File "train.py", line 210, in train
    output,step_coverage_loss = model(texts, summaries_modified,id2elmoid_vectorized,numero_oov,texts_extend_vocab)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/apex/amp/_initialize.py", line 168, in new_fwd
    **applier(kwargs, input_caster))
  File "/content/drive/My Drive/Tesi_Text_Summarization/colab/tesi/pointer_generator.py", line 54, in forward
    encoder_outputs, hidden = self.encoder(texts)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/apex/amp/_initialize.py", line 168, in new_fwd
    **applier(kwargs, input_caster))
  File "/content/drive/My Drive/Tesi_Text_Summarization/colab/tesi/encoder.py", line 22, in forward
    outputs, hidden = self.gru(elmo_embeddings)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/rnn.py", line 179, in forward
    self.dropout, self.training, self.bidirectional, self.batch_first)
RuntimeError: param_from.type() == param_to.type() ASSERT FAILED at /pytorch/aten/src/ATen/native/cudnn/RNN.cpp:539, please report a bug to PyTorch. parameter types mismatch

I think it's due amp rather than PyTorch itself, but let me know what do you think. Also, I'm not sure if in a seq2seq mode I should pass encoder, decoder and attention weights separately to the amp.initalize() or I should pass the the object holding them (in my case the variable model)

Start with train.py and follow the model call into pointer_generator.py and encoder.py, especially the GRU invocation and the amp.initialize() setup described in the issue. Reproduce the traceback, inspect the AMP and RNN parameter types, and document a confirmed cause and the correct initialization scope.