torch2trt: max_batch_size doesn't work as expected

As per the documentation, I should be able to provide param max_batch_size during model conversion and be able to provide different batch sizes during inference. This doesn’t seem to be happening. Code to reproduce:

import torch
from torch2trt import torch2trt
from torchvision.models.alexnet import alexnet

# create some regular pytorch model...
model = alexnet(pretrained=True).eval().cuda()

# create example data
x = torch.ones((1, 3, 224, 224)).cuda()

# convert to TensorRT feeding sample data as input
model_trt = torch2trt(model, [x], max_batch_size=4)

nx = torch.ones((3, 3, 224, 224)).cuda()
y = model(nx)
y_trt = model_trt(nx)
print (f'y shape: {y.shape}, y_trt shape: {y_trt.shape}')

# check the output against PyTorch
print(torch.max(torch.abs(y - y_trt)))

Output:

[05/24/2022-02:31:31] [TRT] [E] 3: [executionContext.cpp::setBindingDimensions::944] Error Code 3: API Usage Error (Parameter check failed at: runtime/api/executionContext.cpp::setBindingDimensions::944, condition: profileMaxDims.d[i] >= dimensions.d[i]. Supplied binding dimension [3,3,224,224] for bindings[0] exceed min ~ max range at index 0, maximum dimension in profile is 1, minimum dimension in profile is 1, but supplied dimension is 3.
)
y shape: torch.Size([3, 1000]), y_trt shape: torch.Size([1, 1000])
tensor(1.9073e-06, device='cuda:0', grad_fn=<MaxBackward1>)

As we can see the converted TRT model’s output batch size is 1 instead of 3. Does max_batch_size not work?