Issue with MTEncDecModel model in ONNX IOBinding
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

I face the following error when exporting [MTEncDecModel](https://github.com/NVIDIA/NeMo/blob/r1.2.0/nemo/collections/nlp/models/machine_translation/mt_enc_dec_model.py) to ONNX and trying to bind inputs and outputs directly to a PyTorch tensor.

**Steps/Code to reproduce bug**

```python
import torch
import numpy as np
import onnxruntime
from nemo.collections.nlp.models import MTEncDecModel

# load the model from pretrained
model = MTEncDecModel.from_pretrained('nmt_en_de_transformer24x6')

# export all model components to onnx:
model.encoder.export('encoder.onnx')
model.decoder.export('decoder.onnx')
model.log_softmax.export('classifier.onnx')

# initialise all the onnx sessions and bindings:
encoder_session = onnxruntime.InferenceSession('encoder.onnx', providers=['CUDAExecutionProvider'])
decoder_session = onnxruntime.InferenceSession('decoder.onnx', providers=['CUDAExecutionProvider'])
classifier_session = onnxruntime.InferenceSession('classifier.onnx', providers=['CUDAExecutionProvider'])

encoder_session_binding = encoder_session.io_binding()
decoder_session_binding = decoder_session.io_binding()
classifier_session_binding = classifier_session.io_binding()

# preprocess the data using the original nemo model for simplicity:
TEXT = ['good morning', 'hello world']
src_ids, src_mask = model.prepare_inference_batch(TEXT)

# bind encoder inputs and outputs to a PyTorch tensor
src_ids_tensor, src_mask_tensor = src_ids.contiguous(), src_mask.contiguous()

encoder_session_binding.bind_input(
    name='input_ids',
    device_type='cuda',
    device_id=0,
    element_type=np.int64,
    shape=tuple(src_ids_tensor.shape),
    buffer_ptr=src_ids_tensor.data_ptr(),
)
encoder_session_binding.bind_input(
    name='encoder_mask',
    device_type='cuda',
    device_id=0,
    element_type=np.int64,
    shape=tuple(src_mask_tensor.shape),
    buffer_ptr=src_mask_tensor.data_ptr(),
)

encoder_hidden_state_shape = (src_ids.shape[0], src_ids.shape[1], 1024)
encoder_hidden_state = torch.empty(encoder_hidden_state_shape, dtype=torch.float32, device='cuda:0').contiguous()
encoder_session_binding.bind_output(
    name='last_hidden_states',
    device_type='cuda',
    device_id=0,
    element_type=np.float32,
    shape=tuple(encoder_hidden_state.shape),
    buffer_ptr=encoder_hidden_state.data_ptr(),
)

encoder_session.run_with_iobinding(encoder_session_binding)

# simple greedy search:
MAX_GENERATION_DELTA = 5
BOS = model.encoder_tokenizer.bos_id
EOS = model.encoder_tokenizer.eos_id
PAD = model.encoder_tokenizer.pad_id

def decode(tgt: torch.tensor, embeding: torch.tensor, src_mask: torch.tensor, decoder_hidden_state: torch.tensor) -> torch.tensor:
    # bind decoder inputs and outputs to a PyTorch tensor
    tgt_tensor = tgt.contiguous()
    embeding_tensor = embeding.contiguous()
    src_mask_tensor = src_mask.contiguous()
    decoder_hidden_state_tensor = decoder_hidden_state.contiguous()
    decoder_mask_tensor = (tgt != PAD).long().contiguous()
    
    decoder_session_binding.bind_input(
        name='input_ids',
        device_type='cuda',
        device_id=0,
        element_type=np.int64,
        shape=tuple(tgt_tensor.shape),
        buffer_ptr=tgt_tensor.data_ptr(),
    )

    decoder_session_binding.bind_input(
        name='decoder_mask',
        device_type='cuda',
        device_id=0,
        element_type=np.int64,
        shape=tuple(decoder_mask_tensor.shape),
        buffer_ptr=decoder_mask_tensor.data_ptr(),
    )

    decoder_session_binding.bind_input(
        name='encoder_mask',
        device_type='cuda',
        device_id=0,
        element_type=np.int64,
        shape=tuple(src_mask_tensor.shape),
        buffer_ptr=src_mask_tensor.data_ptr(),
    )

    decoder_session_binding.bind_input(
        name='encoder_embeddings',
        device_type='cuda',
        device_id=0,
        element_type=np.float32,
        shape=tuple(embeding_tensor.shape),
        buffer_ptr=embeding_tensor.data_ptr(),
    )

    decoder_session_binding.bind_input(
        name='decoder_mems',
        device_type='cuda',
        device_id=0,
        element_type=np.float32,
        shape=tuple(decoder_hidden_state_tensor.shape),
        buffer_ptr=decoder_hidden_state_tensor.data_ptr(),
    )
    
    decoder_hidden_state_shape = (embeding.shape[0], 8, decoder_hidden_state.shape[2]+1, 1024)
    decoder_hidden_state_tensor = torch.empty(decoder_hidden_state_shape, dtype=torch.float32, device='cuda:0').contiguous()
    decoder_session_binding.bind_output(
        name='last_hidden_states',
        device_type='cuda',
        device_id=0,
        element_type=np.float32,
        shape=tuple(decoder_hidden_state_tensor.shape),
        buffer_ptr=decoder_hidden_state_tensor.data_ptr(),
    )
    decoder_session.run_with_iobinding(decoder_session_binding)

    # bind classifier inputs and outputs to a PyTorch tensor
    classifier_session_binding.bind_input(
        name='hidden_states',
        device_type='cuda',
        device_id=0,
        element_type=np.int64,
        shape=tuple(decoder_hidden_state[:,-1,-1:].shape),
        buffer_ptr=decoder_hidden_state[:,-1,-1:].data_ptr(),
    )
    log_probs_shape = (embeding.shape[0], 1, 32000)
    log_probs_tensor = torch.empty(log_probs_shape, dtype=torch.float32, device='cuda:0').contiguous()
    classifier_session_binding.bind_output(
        name='log_probs',
        device_type='cuda',
        device_id=0,
        element_type=np.int64,
        shape=tuple(log_probs_tensor.shape),
        buffer_ptr=log_probs_tensor.data_ptr(),
    )

    classifier_session.run_with_iobinding(classifier_session_binding)

    return log_probs_tensor, decoder_hidden_state_tensor

max_out_len = encoder_hidden_state.shape[1] + MAX_GENERATION_DELTA
tgt=torch.full((encoder_hidden_state.shape[0], max_out_len), PAD)
tgt[:, 0] = BOS
decoder_hidden_state = torch.ones(encoder_hidden_state.shape[0], 8, 0, encoder_hidden_state.shape[-1]).to(torch.float32)

for i in range(1, max_out_len):
    log_probs, decoder_hidden_state = decode(tgt[:, :i], encoder_hidden_state, src_mask, decoder_hidden_state)
    next_tokens = log_probs[:, -1].argmax(axis=1)
    tgt[:, i] = next_tokens
    if ((tgt == EOS).sum(axis=1) > 0).all():
        break
```

**Error:**
```
RuntimeError                              Traceback (most recent call last)
/tmp/ipykernel_8614/803687795.py in <module>
      5 
      6 for i in range(1, max_out_len):
----> 7     log_probs, decoder_hidden_state = decode(tgt[:, :i], encoder_hidden_state, src_mask, decoder_hidden_state)
      8     next_tokens = log_probs[:, -1].argmax(axis=1) # NOTE: ONNX decoder returns multiple outputs which is different to pytorch version, so I get the last one (this could be where error is?)
      9     tgt[:, i] = next_tokens

/tmp/ipykernel_8614/2263200407.py in decode(tgt, embeding, src_mask, decoder_hidden_state)
     42     )
     43 
---> 44     decoder_session_binding.bind_input(
     45         name='decoder_mems',
     46         device_type='cuda',

/opt/conda/lib/python3.8/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py in bind_input(self, name, device_type, device_id, element_type, shape, buffer_ptr)
    439         :param buffer_ptr: memory pointer to input data
    440         '''
--> 441         self._iobinding.bind_input(name,
    442                                    C.OrtDevice(get_ort_device_type(device_type), C.OrtDevice.default_memory(),
    443                                                device_id),

RuntimeError: /onnxruntime_src/onnxruntime/python/onnxruntime_pybind_iobinding.cc:64 onnxruntime::python::addIoBindingMethods(pybind11::module&)::<lambda(onnxruntime::SessionIOBinding*, const string&, const OrtDevice&, pybind11::object&, const std::vector<long int>&, int64_t)> data_ptr != 0 was false. Pointer to data memory is not valid
```

**Environment overview (please complete the following information)**

 - Environment location: Docker on Kubernetes
 - Method of NeMo install: using helm chart using nvcr.io/nvidia/nemo: 1.6.1
 - Additionally install `pip install onnxruntime-gpu`