How can I use bfloat16 with tensorrt?

[vedantroy]

I'm trying to use torch.compile as follows:

        import torch_tensorrt
        vae.encoder = torch.compile(vae.encoder, mode="reduce-overhead", fullgraph=True, backend="tensorrt")

But i'm getting the error:

  File "/tmp/train_env_v017/lib/python3.10/site-packages/torch_tensorrt/_enums.py", line 116, in _from
    elif t == trt.bf16:
torch._dynamo.exc.BackendCompilerFailed: backend='tensorrt' raised:
AttributeError: module 'tensorrt' has no attribute 'bf16'

While executing %clone_1 : [num_users=1] = call_function[target=torch.ops.aten.clone.default](args = (%expand,), kwargs = {memory_format: torch.contiguous_format, _itensor_to_tensor_meta: {<tensorrt_bindings.tensorrt.ITensor object at 0x148576b70130>: ((6, 37, 192, 336, 15), torch.bfloat16, False, (35804160, 64512, 336, 1, 2386944), None, False, {}), <tensorrt_bindings.tensorrt.ITensor object at 0x148576b701b0>: ((6, 37, 192, 336, 15), torch.bfloat16, False, (35804160, 64512, 336, 1, 2386944), None, False, {})}})

Do I need to be using a very recent version of tensorrt? My current versions are:

tensorrt                      10.0.1
tensorrt-cu12                 10.0.1
tensorrt-cu12-bindings        10.0.1
tensorrt-cu12-libs            10.0.1
torch                         2.4.0.dev20240522+cu121
torch-fidelity                0.3.0
torch-optimizer               0.3.0
torch-summary                 1.4.5
torch_tensorrt                2.4.0.dev20240522+cu121
torchaudio                    2.2.0.dev20240522+cu121
torchdiffeq                   0.2.3
torchmetrics                  1.3.0.post0
torchvision                   0.19.0.dev20240522+cu121

I wonder if this is a bug, because replacing trt.bf16 in the code with trt.bfloat16 makes the crash go-away, although then it gets replaced with this:

#2902

[narendasan]

The original issue has been fixed. The other one will be.