Closed pngmafia closed 1 year ago
Can you provide a way to reproduce the exception?
I just re ran the cells on the gpt2 notebook and this issue occurs while running the tensorrt inference cell
tensorrt_model: Callable[[Dict[str, torch.Tensor]], torch.Tensor] = load_engine(
engine_file_path="test-gpt2.plan", runtime=runtime
)
def inference_tensorrt(input_ids: torch.Tensor) -> torch.Tensor:
data = {"input_ids": input_ids}
return tensorrt_model(data)
gpt2_model = GPTModelWrapper(config=model.config, device=torch.device("cuda"), inference=inference_tensorrt)
inputs.to("cuda")
sample_output = gpt2_model.generate(inputs.input_ids, max_length=64)
print(tokenizer.decode(sample_output[0], skip_special_tokens=True))
for _ in range(2):
_ = gpt2_model.generate(inputs.input_ids, max_length=64)
start = time.time()
for _ in range(10):
_ = gpt2_model.generate(inputs.input_ids, max_length=256)
print(f"----\nTensorRT + CUDA tensors: {(time.time() - start)/10:.2f}/sequence")
del tensorrt_model
Any clue how I could fix this?